Example-01: Optimizaion (import/export API)

[1]:
# Loading API facilitates (suboptimal) interface with different optimization libraries
# In this example, quadrupole gradients are used to fit beta functions
[2]:
from pathlib import Path

from model.command.external import load_sdds
from model.command.external import load_lattice
from model.command.external import text_lattice

import numpy
from numpy import ndarray as Array

from scipy.optimize import minimize
[3]:
# Set quadrupole gradient and compute and return twiss parameters

def evaluate(knobs:Array) -> Array:
    kf, kd = knobs
    path:Path = Path('optimize.lte')
    lattice:dict[str, dict[str, str | int | float | dict]] = load_lattice(path)
    lattice['QF']['K1'] = float(kf)
    lattice['QD']['K1'] = float(kd)
    with path.open('w') as stream:
        stream.write(text_lattice('LTE', lattice))
    !elegant 'optimize.ele' > /dev/null
    !sddsconvert -ascii 'binary.twiss' 'optimize.twiss'
    path:Path = Path('optimize.twiss')
    _, columns = load_sdds(path)
    return numpy.asarray([[data['betax'], data['betay']] for location, data in columns.items()]).T
[4]:
# Set target beta functions

target:Array = numpy.asarray([+0.21, -0.19])
result:Array = evaluate(target)
[5]:
# Set objetive function to minimize

def objective(knobs:Array) -> Array:
    return numpy.sum((evaluate(knobs) - result)**2)

objective(target)
[5]:
0.0
[6]:
# Optimize

knobs:Array = numpy.asarray([+0.20, -0.20])

minimize(objective, knobs, method='Nelder-Mead')
[6]:
       message: Optimization terminated successfully.
       success: True
        status: 0
           fun: 1.3489521479279302e-28
             x: [ 2.100e-01 -1.900e-01]
           nit: 20
          nfev: 40
 final_simplex: (array([[ 2.100e-01, -1.900e-01],
                       [ 2.100e-01, -1.900e-01],
                       [ 2.100e-01, -1.900e-01]]), array([ 1.349e-28,  2.891e-05,  6.010e-05]))

Example-02: Workflow (MADX)

[1]:
from pathlib import Path

from model.command.external import load_lattice
from model.command.external import rift_lattice
from model.command.external import text_lattice
from model.command.external import load_tfs
from model.command.external import convert
from model.command.external import add_rc
[2]:
# Given some initial MADX lattice file (FODO)

file = Path('initial.madx')

with file.open('r') as stream:
    print(stream.read())

# Several regular elements are defined
# HEAD and TAIL should appear as the first and the last elements
# All elements should be defined on a single line with numerical parameters
# Lattice should be defined using lines
# Comma after element type is mandatory
# Comments appearing after definitions should also represent an element definition
DR: DRIFT, L=2.0;
BM: SBEND, L=1.0, ANGLE=0.17453292519943295;
QF: QUADRUPOLE, L=1.0, K1=+0.2;
QD: QUADRUPOLE, L=0.5, K1=-0.2;

M: MONITOR,;

HEAD: MARKER,; ! TEST: DRIFT,
TAIL: MARKER,; ! TEST: DRIFT,

FODO: LINE=(HEAD, M, QD, DR, BM, DR, QF, DR, BM, DR, QD, TAIL) ;

[3]:
# If element and beamline definitions comply with the above requirements
# The lattice file can be loaded as a python dictionary

lattice = load_lattice(file)

for key, value in lattice.items():
    print(key, value)

# For each element and beamline, a key-value pair in created
# Value is itself a dictionary containing all information about the original elements
# Each element parameter is casted from string to int, float or string
# Comment after element definition is saved into RC (it has a special use case, see below)
DR {'KIND': 'DRIFT', 'RC': '', 'L': 2.0}
BM {'KIND': 'SBEND', 'RC': '', 'L': 1.0, 'ANGLE': 0.17453292519943295}
QF {'KIND': 'QUADRUPOLE', 'RC': '', 'L': 1.0, 'K1': 0.2}
QD {'KIND': 'QUADRUPOLE', 'RC': '', 'L': 0.5, 'K1': -0.2}
M {'KIND': 'MONITOR', 'RC': ''}
HEAD {'KIND': 'MARKER', 'RC': 'TEST: DRIFT,'}
TAIL {'KIND': 'MARKER', 'RC': 'TEST: DRIFT,'}
FODO {'KIND': 'LINE', 'SEQUENCE': ['HEAD', 'M', 'QD', 'DR', 'BM', 'DR', 'QF', 'DR', 'BM', 'DR', 'QD', 'TAIL']}
[4]:
# Error lattice is defined by a set of linear transformations between selected locations
# Each locations can be a MONITOR (beam observation) or a VIRTUAL (error)
# Two special locatons (HEAD and TAIL) should present in the lattice
# Using the above dictionary representation, new observation locations can be inserted
# Locations are inserted at the middle of selected elements (selected by type or name)
# Selected elements are splitted in half and renamed, old names are binded to beamlines
# Original element definitions are added to created location RC
# Typicaly, monitor locations correspond to MONITOR elements, but new monitor elements can be also inserted
# Virtual locations can be inserted into quadrupole or other elements to represent errors

lattice = rift_lattice(lattice,
                       'MONITOR',
                       'MARKER',
                       ['DRIFT'],
                       ['SBEND', 'QUADRUPOLE'],
                       exclude_virtual=['QD'])

for key, value in lattice.items():
    print(key, value)
M_DR {'KIND': 'MONITOR', 'RC': ['DR', {'KIND': 'DRIFT', 'L': 2.0}]}
H_DR {'KIND': 'DRIFT', 'L': 1.0}
DR {'KIND': 'LINE', 'SEQUENCE': ['H_DR', 'M_DR', 'H_DR']}
V_BM {'KIND': 'MARKER', 'RC': ['BM', {'KIND': 'SBEND', 'L': 1.0, 'ANGLE': 0.17453292519943295}]}
H_BM {'KIND': 'SBEND', 'L': 0.5, 'ANGLE': 0.08726646259971647}
BM {'KIND': 'LINE', 'SEQUENCE': ['H_BM', 'V_BM', 'H_BM']}
V_QF {'KIND': 'MARKER', 'RC': ['QF', {'KIND': 'QUADRUPOLE', 'L': 1.0, 'K1': 0.2}]}
H_QF {'KIND': 'QUADRUPOLE', 'L': 0.5, 'K1': 0.2}
QF {'KIND': 'LINE', 'SEQUENCE': ['H_QF', 'V_QF', 'H_QF']}
QD {'KIND': 'QUADRUPOLE', 'RC': '', 'L': 0.5, 'K1': -0.2}
M {'KIND': 'MONITOR', 'RC': ''}
HEAD {'KIND': 'MARKER', 'RC': 'TEST: DRIFT,'}
TAIL {'KIND': 'MARKER', 'RC': 'TEST: DRIFT,'}
FODO {'KIND': 'LINE', 'SEQUENCE': ['HEAD', 'M', 'QD', 'DR', 'BM', 'DR', 'QF', 'DR', 'BM', 'DR', 'QD', 'TAIL']}
[5]:
# Modified lattice can be converted to text
# Comments are added to locations while original comments are preserved

text = text_lattice('MADX', lattice, rc=True)

print(text)
M_DR: MONITOR,; ! DR: DRIFT, L=2.0,;
H_DR: DRIFT, L=1.0,;
DR: LINE=(H_DR, M_DR, H_DR);
V_BM: MARKER,; ! BM: SBEND, L=1.0, ANGLE=0.17453292519943295,;
H_BM: SBEND, L=0.5, ANGLE=0.08726646259971647,;
BM: LINE=(H_BM, V_BM, H_BM);
V_QF: MARKER,; ! QF: QUADRUPOLE, L=1.0, K1=0.2,;
H_QF: QUADRUPOLE, L=0.5, K1=0.2,;
QF: LINE=(H_QF, V_QF, H_QF);
QD: QUADRUPOLE, L=0.5, K1=-0.2,;
M: MONITOR,;
HEAD: MARKER,; ! TEST: DRIFT,
TAIL: MARKER,; ! TEST: DRIFT,
FODO: LINE=(HEAD, M, QD, DR, BM, DR, QF, DR, BM, DR, QD, TAIL);

[6]:
# Compute TWISS parameters using MADX
# TWISS command is appended to modified lattice

task = """
BEAM;
USE, PERIOD=FODO;
SET,FORMAT="20.20f","-20s";
TWISS;
WRITE,TABLE=TWISS,FILE="final.tfs";
RETURN;
""" ;

with Path('final.madx').open('w') as stream:
    stream.write(text)
    stream.write(task)

!madx final.madx > /dev/null
[7]:
# Load lattice can be also loaded from file
# Original comments will be parsed as elements (look at HEAD and TAIL)
# Empty RC will be nested in this case

file = Path('final.madx')

with file.open('w') as stream:
    stream.write(text)

lattice = load_lattice(file, rc=True)

for key, value in lattice.items():
    print(key, value)
M_DR {'KIND': 'MONITOR', 'RC': ['DR', {'KIND': 'DRIFT', 'RC': '', 'L': 2.0}]}
H_DR {'KIND': 'DRIFT', 'RC': ['', {'KIND': '', 'RC': ''}], 'L': 1.0}
DR {'KIND': 'LINE', 'SEQUENCE': ['H_DR', 'M_DR', 'H_DR']}
V_BM {'KIND': 'MARKER', 'RC': ['BM', {'KIND': 'SBEND', 'RC': '', 'L': 1.0, 'ANGLE': 0.17453292519943295}]}
H_BM {'KIND': 'SBEND', 'RC': ['', {'KIND': '', 'RC': ''}], 'L': 0.5, 'ANGLE': 0.08726646259971647}
BM {'KIND': 'LINE', 'SEQUENCE': ['H_BM', 'V_BM', 'H_BM']}
V_QF {'KIND': 'MARKER', 'RC': ['QF', {'KIND': 'QUADRUPOLE', 'RC': '', 'L': 1.0, 'K1': 0.2}]}
H_QF {'KIND': 'QUADRUPOLE', 'RC': ['', {'KIND': '', 'RC': ''}], 'L': 0.5, 'K1': 0.2}
QF {'KIND': 'LINE', 'SEQUENCE': ['H_QF', 'V_QF', 'H_QF']}
QD {'KIND': 'QUADRUPOLE', 'RC': ['', {'KIND': '', 'RC': ''}], 'L': 0.5, 'K1': -0.2}
M {'KIND': 'MONITOR', 'RC': ['', {'KIND': '', 'RC': ''}]}
HEAD {'KIND': 'MARKER', 'RC': ['TEST', {'KIND': 'DRIFT', 'RC': ''}]}
TAIL {'KIND': 'MARKER', 'RC': ['TEST', {'KIND': 'DRIFT', 'RC': ''}]}
FODO {'KIND': 'LINE', 'SEQUENCE': ['HEAD', 'M', 'QD', 'DR', 'BM', 'DR', 'QF', 'DR', 'BM', 'DR', 'QD', 'TAIL']}
[8]:
# TWISS results can be loaded into python dictionaries

data = Path('final.tfs')
parameters, columns = load_tfs(data)
[9]:
# Optics data can be converted into model table
# Note, all locations have different name
# If an element appear several times in a line, locations are renamed

table = convert(columns, 'TFS', ['MONITOR'], ['MARKER'], rc=True)
table
[9]:
{'HEAD': {'TYPE': 'VIRTUAL',
  'S': 0.0,
  'BX': 4.287017735718936,
  'AX': -3.338e-16,
  'FX': 0.0,
  'BY': 19.818489282044894,
  'AY': -2.975e-17,
  'FY': 0.0,
  'DQX': 1.5490410441348796,
  'DPX': 5.551e-17,
  'DQY': 0.0,
  'DPY': -0.0,
  'RC': None},
 'M': {'TYPE': 'MONITOR',
  'S': 0.0,
  'BX': 4.287017735718936,
  'AX': -3.338e-16,
  'FX': 0.0,
  'BY': 19.818489282044894,
  'AY': -2.975e-17,
  'FY': 0.0,
  'DQX': 1.5490410441348796,
  'DPX': 5.551e-17,
  'DQY': 0.0,
  'DPY': -0.0,
  'RC': None},
 'M_DR': {'TYPE': 'MONITOR',
  'S': 1.5,
  'BX': 5.980356480296974,
  'AX': -0.8524040348462865,
  'FX': 0.3068185833848285,
  'BY': 15.315159900296642,
  'AY': 1.6491678474974667,
  'FY': 0.08453149247893033,
  'DQX': 1.744126900814982,
  'DPX': 0.1561982029636459,
  'DQY': 0.0,
  'DPY': 0.0,
  'RC': None},
 'V_BM': {'TYPE': 'VIRTUAL',
  'S': 3.0,
  'BX': 9.120629409698159,
  'AX': -1.1465687400023221,
  'FX': 0.5107301354647854,
  'BY': 10.914137614299745,
  'AY': 1.2848470098337954,
  'FY': 0.20081415045191314,
  'DQX': 1.992896582518225,
  'DPX': 0.21385269164968151,
  'DQY': 0.0,
  'DPY': 0.0,
  'RC': None},
 'M_DR_1': {'TYPE': 'MONITOR',
  'S': 4.5,
  'BX': 12.70896979457465,
  'AX': -1.3363966727488819,
  'FX': 0.6500917125372264,
  'BY': 7.606077841293864,
  'AY': 0.9205261721701236,
  'FY': 0.3661997931359615,
  'DQX': 2.3837861006470566,
  'DPX': 0.26987963222618605,
  'DQY': 0.0,
  'DPY': 0.0,
  'RC': None},
 'V_QF': {'TYPE': 'VIRTUAL',
  'S': 6.0,
  'BX': 16.392007194825762,
  'AX': -8.0415e-16,
  'FX': 0.7521811118347426,
  'BY': 5.674619594695141,
  'AY': -2.5871e-16,
  'FY': 0.601131166393595,
  'DQX': 2.7214181783177667,
  'DPX': 5.551e-17,
  'DQY': 0.0,
  'DPY': 0.0,
  'RC': None},
 'M_DR_2': {'TYPE': 'MONITOR',
  'S': 7.5,
  'BX': 12.708969794574653,
  'AX': 1.336396672748881,
  'FX': 0.8542705111322586,
  'BY': 7.606077841293868,
  'AY': -0.9205261721701244,
  'FY': 0.8360625396512283,
  'DQX': 2.3837861006470566,
  'DPX': -0.26987963222618594,
  'DQY': 0.0,
  'DPY': 0.0,
  'RC': None},
 'V_BM_1': {'TYPE': 'VIRTUAL',
  'S': 9.0,
  'BX': 9.120629409698164,
  'AX': 1.146568740002322,
  'FX': 0.9936320882046995,
  'BY': 10.914137614299747,
  'AY': -1.2848470098337954,
  'FY': 1.0014481823352765,
  'DQX': 1.9928965825182252,
  'DPX': -0.21385269164968146,
  'DQY': 0.0,
  'DPY': 0.0,
  'RC': None},
 'M_DR_3': {'TYPE': 'MONITOR',
  'S': 10.5,
  'BX': 5.980356480296976,
  'AX': 0.8524040348462865,
  'FX': 1.1975436402846564,
  'BY': 15.315159900296639,
  'AY': -1.6491678474974665,
  'FY': 1.1177308403082595,
  'DQX': 1.7441269008149822,
  'DPX': -0.1561982029636459,
  'DQY': 0.0,
  'DPY': 0.0,
  'RC': None},
 'TAIL': {'TYPE': 'VIRTUAL',
  'S': 12.0,
  'BX': 4.287017735718939,
  'AX': 3.6486e-16,
  'FX': 1.5043622236694847,
  'BY': 19.81848928204488,
  'AY': -5.5051e-16,
  'FY': 1.2022623327871897,
  'DQX': 1.54904104413488,
  'DPX': -2.776e-17,
  'DQY': 0.0,
  'DPY': 0.0,
  'RC': None}}
[10]:
# RC parameter from lattice data can be added to model table
# Configuration table can be saved using util.save

table = add_rc(table, lattice)
table
[10]:
{'HEAD': {'TYPE': 'VIRTUAL',
  'S': 0.0,
  'BX': 4.287017735718936,
  'AX': -3.338e-16,
  'FX': 0.0,
  'BY': 19.818489282044894,
  'AY': -2.975e-17,
  'FY': 0.0,
  'DQX': 1.5490410441348796,
  'DPX': 5.551e-17,
  'DQY': 0.0,
  'DPY': -0.0,
  'RC': None},
 'M': {'TYPE': 'MONITOR',
  'S': 0.0,
  'BX': 4.287017735718936,
  'AX': -3.338e-16,
  'FX': 0.0,
  'BY': 19.818489282044894,
  'AY': -2.975e-17,
  'FY': 0.0,
  'DQX': 1.5490410441348796,
  'DPX': 5.551e-17,
  'DQY': 0.0,
  'DPY': -0.0,
  'RC': None},
 'M_DR': {'TYPE': 'MONITOR',
  'S': 1.5,
  'BX': 5.980356480296974,
  'AX': -0.8524040348462865,
  'FX': 0.3068185833848285,
  'BY': 15.315159900296642,
  'AY': 1.6491678474974667,
  'FY': 0.08453149247893033,
  'DQX': 1.744126900814982,
  'DPX': 0.1561982029636459,
  'DQY': 0.0,
  'DPY': 0.0,
  'RC': ['DR', {'KIND': 'DRIFT', 'L': 2.0}]},
 'V_BM': {'TYPE': 'VIRTUAL',
  'S': 3.0,
  'BX': 9.120629409698159,
  'AX': -1.1465687400023221,
  'FX': 0.5107301354647854,
  'BY': 10.914137614299745,
  'AY': 1.2848470098337954,
  'FY': 0.20081415045191314,
  'DQX': 1.992896582518225,
  'DPX': 0.21385269164968151,
  'DQY': 0.0,
  'DPY': 0.0,
  'RC': ['BM', {'KIND': 'SBEND', 'L': 1.0, 'ANGLE': 0.17453292519943295}]},
 'M_DR_1': {'TYPE': 'MONITOR',
  'S': 4.5,
  'BX': 12.70896979457465,
  'AX': -1.3363966727488819,
  'FX': 0.6500917125372264,
  'BY': 7.606077841293864,
  'AY': 0.9205261721701236,
  'FY': 0.3661997931359615,
  'DQX': 2.3837861006470566,
  'DPX': 0.26987963222618605,
  'DQY': 0.0,
  'DPY': 0.0,
  'RC': ['DR', {'KIND': 'DRIFT', 'L': 2.0}]},
 'V_QF': {'TYPE': 'VIRTUAL',
  'S': 6.0,
  'BX': 16.392007194825762,
  'AX': -8.0415e-16,
  'FX': 0.7521811118347426,
  'BY': 5.674619594695141,
  'AY': -2.5871e-16,
  'FY': 0.601131166393595,
  'DQX': 2.7214181783177667,
  'DPX': 5.551e-17,
  'DQY': 0.0,
  'DPY': 0.0,
  'RC': ['QF', {'KIND': 'QUADRUPOLE', 'L': 1.0, 'K1': 0.2}]},
 'M_DR_2': {'TYPE': 'MONITOR',
  'S': 7.5,
  'BX': 12.708969794574653,
  'AX': 1.336396672748881,
  'FX': 0.8542705111322586,
  'BY': 7.606077841293868,
  'AY': -0.9205261721701244,
  'FY': 0.8360625396512283,
  'DQX': 2.3837861006470566,
  'DPX': -0.26987963222618594,
  'DQY': 0.0,
  'DPY': 0.0,
  'RC': ['DR', {'KIND': 'DRIFT', 'L': 2.0}]},
 'V_BM_1': {'TYPE': 'VIRTUAL',
  'S': 9.0,
  'BX': 9.120629409698164,
  'AX': 1.146568740002322,
  'FX': 0.9936320882046995,
  'BY': 10.914137614299747,
  'AY': -1.2848470098337954,
  'FY': 1.0014481823352765,
  'DQX': 1.9928965825182252,
  'DPX': -0.21385269164968146,
  'DQY': 0.0,
  'DPY': 0.0,
  'RC': ['BM', {'KIND': 'SBEND', 'L': 1.0, 'ANGLE': 0.17453292519943295}]},
 'M_DR_3': {'TYPE': 'MONITOR',
  'S': 10.5,
  'BX': 5.980356480296976,
  'AX': 0.8524040348462865,
  'FX': 1.1975436402846564,
  'BY': 15.315159900296639,
  'AY': -1.6491678474974665,
  'FY': 1.1177308403082595,
  'DQX': 1.7441269008149822,
  'DPX': -0.1561982029636459,
  'DQY': 0.0,
  'DPY': 0.0,
  'RC': ['DR', {'KIND': 'DRIFT', 'L': 2.0}]},
 'TAIL': {'TYPE': 'VIRTUAL',
  'S': 12.0,
  'BX': 4.287017735718939,
  'AX': 3.6486e-16,
  'FX': 1.5043622236694847,
  'BY': 19.81848928204488,
  'AY': -5.5051e-16,
  'FY': 1.2022623327871897,
  'DQX': 1.54904104413488,
  'DPX': -2.776e-17,
  'DQY': 0.0,
  'DPY': 0.0,
  'RC': None}}

Example-03: Workflow (ELEGANT)

[1]:
from pathlib import Path

from model.command.external import load_lattice
from model.command.external import rift_lattice
from model.command.external import text_lattice
from model.command.external import load_sdds
from model.command.external import convert
from model.command.external import add_rc
[2]:
# Given some initial ELEGANT lattice file (FODO)

file = Path('initial.lte')

with file.open('r') as stream:
    print(stream.read())

# Several regular elements are defined
# HEAD and TAIL should appear as the first and the last elements
# All elements should be defined on a single line with numerical parameters
# Lattice should be defined using lines
# Comma after element type is mandatory
# Comments appearing after definitions should also represent an element definition
DR: DRIF,l=2

BM: CSBEND,angle=0.17453292519943295,fint=0,l=1.0
QD: QUAD,k1=-0.2,l=0.5
QF: QUAD,k1=0.2,l=1.0

M: MONI,

HEAD: MARK, ! TEST: DRIF,
TAIL: MARK, ! TEST: DRIF,

FODO: LINE=(HEAD, M, QD, DR, BM, DR, QF, DR, BM, DR, QD, TAIL)
[3]:
# If element and beamline definitions comply with the above requirements
# The lattice file can be loaded as a python dictionary

lattice = load_lattice(file)

for key, value in lattice.items():
    print(key, value)

# For each element and beamline, a key-value pair in created
# Value is itself a dictionary containing all information about the original elements
# Each element parameter is casted from string to int, float or string
# Comment after element definition is saved into RC (it has a special use case, see below)
DR {'KIND': 'DRIF', 'RC': '', 'L': 2}
BM {'KIND': 'CSBEND', 'RC': '', 'ANGLE': 0.17453292519943295, 'FINT': 0, 'L': 1.0}
QD {'KIND': 'QUAD', 'RC': '', 'K1': -0.2, 'L': 0.5}
QF {'KIND': 'QUAD', 'RC': '', 'K1': 0.2, 'L': 1.0}
M {'KIND': 'MONI', 'RC': ''}
HEAD {'KIND': 'MARK', 'RC': 'TEST: DRIF,'}
TAIL {'KIND': 'MARK', 'RC': 'TEST: DRIF,'}
FODO {'KIND': 'LINE', 'SEQUENCE': ['HEAD', 'M', 'QD', 'DR', 'BM', 'DR', 'QF', 'DR', 'BM', 'DR', 'QD', 'TAIL']}
[4]:
# Error lattice is defined by a set of linear transformations between selected locations
# Each locations can be a MONITOR (beam observation) or a VIRTUAL (error)
# Two special locatons (HEAD and TAIL) should present in the lattice
# Using the above dictionary representation, new observation locations can be inserted
# Locations are inserted at the middle of selected elements (selected by type or name)
# Selected elements are splitted in half and renamed, old names are binded to beamlines
# Original element definitions are added to created location RC
# Typicaly, monitor locations correspond to MONITOR elements, but new monitor elements can be also inserted
# Virtual locations can be inserted into quadrupole or other elements to represent errors

lattice = rift_lattice(lattice,
                       'MONI',
                       'MARK',
                       ['DRIF'],
                       ['CSBEND', 'QUAD'],
                       exclude_virtual=['QD'])

for key, value in lattice.items():
    print(key, value)
M_DR {'KIND': 'MONI', 'RC': ['DR', {'KIND': 'DRIF', 'L': 2}]}
H_DR {'KIND': 'DRIF', 'L': 1.0}
DR {'KIND': 'LINE', 'SEQUENCE': ['H_DR', 'M_DR', 'H_DR']}
V_BM {'KIND': 'MARK', 'RC': ['BM', {'KIND': 'CSBEND', 'ANGLE': 0.17453292519943295, 'FINT': 0, 'L': 1.0}]}
H_BM {'KIND': 'CSBEND', 'ANGLE': 0.08726646259971647, 'FINT': 0, 'L': 0.5}
BM {'KIND': 'LINE', 'SEQUENCE': ['H_BM', 'V_BM', 'H_BM']}
QD {'KIND': 'QUAD', 'RC': '', 'K1': -0.2, 'L': 0.5}
V_QF {'KIND': 'MARK', 'RC': ['QF', {'KIND': 'QUAD', 'K1': 0.2, 'L': 1.0}]}
H_QF {'KIND': 'QUAD', 'K1': 0.2, 'L': 0.5}
QF {'KIND': 'LINE', 'SEQUENCE': ['H_QF', 'V_QF', 'H_QF']}
M {'KIND': 'MONI', 'RC': ''}
HEAD {'KIND': 'MARK', 'RC': 'TEST: DRIF,'}
TAIL {'KIND': 'MARK', 'RC': 'TEST: DRIF,'}
FODO {'KIND': 'LINE', 'SEQUENCE': ['HEAD', 'M', 'QD', 'DR', 'BM', 'DR', 'QF', 'DR', 'BM', 'DR', 'QD', 'TAIL']}
[5]:
# Modified lattice can be converted to text
# Comments are added to locations while original comments are preserved

text = text_lattice('LTE', lattice, rc=True)

print(text)
M_DR: MONI, ! DR: DRIF, L=2,
H_DR: DRIF, L=1.0,
DR: LINE=(H_DR, M_DR, H_DR)
V_BM: MARK, ! BM: CSBEND, ANGLE=0.17453292519943295, FINT=0, L=1.0,
H_BM: CSBEND, ANGLE=0.08726646259971647, FINT=0, L=0.5,
BM: LINE=(H_BM, V_BM, H_BM)
QD: QUAD, K1=-0.2, L=0.5,
V_QF: MARK, ! QF: QUAD, K1=0.2, L=1.0,
H_QF: QUAD, K1=0.2, L=0.5,
QF: LINE=(H_QF, V_QF, H_QF)
M: MONI,
HEAD: MARK, ! TEST: DRIF,
TAIL: MARK, ! TEST: DRIF,
FODO: LINE=(HEAD, M, QD, DR, BM, DR, QF, DR, BM, DR, QD, TAIL)

[6]:
# Compute TWISS parameters using ELEGANT
# Separate command file is created

with Path('final.lte').open('w') as stream:
    stream.write(text)

task = """
&run_setup
  use_beamline="FODO",
  lattice = "final.lte",
  p_central_mev = 1000
&end

&run_control
&end

&twiss_output
  filename = "binary.twiss",
  output_at_each_step = 1
&end

&bunched_beam
&end

&track
&end
""" ;

with Path('final.ele').open('w') as stream:
    stream.write(task)

!elegant final.ele > /dev/null
!sddsconvert -ascii binary.twiss final.twiss
[7]:
# Load lattice can be also loaded from file
# Original comments will be parsed as elements (look at HEAD and TAIL)
# Empty RC will be nested in this case

file = Path('final.lte')

with file.open('w') as stream:
    stream.write(text)

lattice = load_lattice(file, rc=True)

for key, value in lattice.items():
    print(key, value)
M_DR {'KIND': 'MONI', 'RC': ['DR', {'KIND': 'DRIF', 'RC': '', 'L': 2}]}
H_DR {'KIND': 'DRIF', 'RC': ['', {'KIND': '', 'RC': ''}], 'L': 1.0}
DR {'KIND': 'LINE', 'SEQUENCE': ['H_DR', 'M_DR', 'H_DR']}
V_BM {'KIND': 'MARK', 'RC': ['BM', {'KIND': 'CSBEND', 'RC': '', 'ANGLE': 0.17453292519943295, 'FINT': 0, 'L': 1.0}]}
H_BM {'KIND': 'CSBEND', 'RC': ['', {'KIND': '', 'RC': ''}], 'ANGLE': 0.08726646259971647, 'FINT': 0, 'L': 0.5}
BM {'KIND': 'LINE', 'SEQUENCE': ['H_BM', 'V_BM', 'H_BM']}
QD {'KIND': 'QUAD', 'RC': ['', {'KIND': '', 'RC': ''}], 'K1': -0.2, 'L': 0.5}
V_QF {'KIND': 'MARK', 'RC': ['QF', {'KIND': 'QUAD', 'RC': '', 'K1': 0.2, 'L': 1.0}]}
H_QF {'KIND': 'QUAD', 'RC': ['', {'KIND': '', 'RC': ''}], 'K1': 0.2, 'L': 0.5}
QF {'KIND': 'LINE', 'SEQUENCE': ['H_QF', 'V_QF', 'H_QF']}
M {'KIND': 'MONI', 'RC': ['', {'KIND': '', 'RC': ''}]}
HEAD {'KIND': 'MARK', 'RC': ['TEST', {'KIND': 'DRIF', 'RC': ''}]}
TAIL {'KIND': 'MARK', 'RC': ['TEST', {'KIND': 'DRIF', 'RC': ''}]}
FODO {'KIND': 'LINE', 'SEQUENCE': ['HEAD', 'M', 'QD', 'DR', 'BM', 'DR', 'QF', 'DR', 'BM', 'DR', 'QD', 'TAIL']}
[8]:
# TWISS results can be loaded into python dictionaries

data = Path('final.twiss')
parameters, columns = load_sdds(data)
[9]:
# Optics data can be converted into model table
# Note, all locations have different name
# If an element appear several times in a line, locations are renamed

table = convert(columns, 'SDDS', ['MONI'], ['MARK'], rc=True)
table
[9]:
{'HEAD': {'S': 0.0,
  'BX': 4.287017734831204,
  'AX': -1.321304551729011e-16,
  'FX': 0.0,
  'DQX': 1.549040841795901,
  'DPX': 2.775557561562891e-17,
  'BY': 19.81848926815186,
  'AY': 6.545730027929358e-16,
  'FY': 0.0,
  'DQY': 0.0,
  'DPY': 0.0,
  'TYPE': 'VIRTUAL',
  'RC': None},
 'M': {'S': 0.0,
  'BX': 4.287017734831204,
  'AX': -1.321304551729011e-16,
  'FX': 0.0,
  'DQX': 1.549040841795901,
  'DPX': 2.775557561562891e-17,
  'BY': 19.81848926815186,
  'AY': 6.545730027929358e-16,
  'FY': 0.0,
  'DQY': 0.0,
  'DPY': 0.0,
  'TYPE': 'MONITOR',
  'RC': None},
 'M_DR': {'S': 1.5,
  'BX': 5.980356479284525,
  'AX': -0.8524040348212205,
  'FX': 0.3068185834444499,
  'DQX': 1.744126672993481,
  'DPX': 0.1561981825607101,
  'BY': 15.31515988971357,
  'AY': 1.649167846239909,
  'FY': 0.08453149253790615,
  'DQY': 0.0,
  'DPY': 0.0,
  'TYPE': 'MONITOR',
  'RC': None},
 'V_BM': {'S': 3.0,
  'BX': 9.120629409314306,
  'AX': -1.146568740048648,
  'FX': 0.5107301355818847,
  'DQX': 1.99289632225274,
  'DPX': 0.2138526637243729,
  'BY': 10.91413760701325,
  'AY': 1.284847009200138,
  'FY': 0.2008141505892614,
  'DQY': 0.0,
  'DPY': 0.0,
  'TYPE': 'VIRTUAL',
  'RC': None},
 'M_DR_1': {'S': 4.5,
  'BX': 12.70896979491113,
  'AX': -1.336396672789745,
  'FX': 0.6500917126795847,
  'DQX': 2.383785789407781,
  'DPX': 0.2698795969893243,
  'BY': 7.606077834992553,
  'AY': 0.9205261718281931,
  'FY': 0.3661997933967667,
  'DQY': 0.0,
  'DPY': 0.0,
  'TYPE': 'MONITOR',
  'RC': None},
 'V_QF': {'S': 6.0,
  'BX': 16.39200719526896,
  'AX': 3.608224830031759e-16,
  'FX': 0.7521811119743597,
  'DQX': 2.721417822995524,
  'DPX': -5.551115123125783e-17,
  'BY': 5.674619589422572,
  'AY': 8.326672684688672e-16,
  'FY': 0.6011311668647427,
  'DQY': 0.0,
  'DPY': 0.0,
  'TYPE': 'VIRTUAL',
  'RC': None},
 'M_DR_2': {'S': 7.5,
  'BX': 12.70896979491113,
  'AX': 1.336396672789746,
  'FX': 0.8542705112691347,
  'DQX': 2.383785789407781,
  'DPX': -0.2698795969893243,
  'BY': 7.606077834992548,
  'AY': -0.920526171828191,
  'FY': 0.8360625403327188,
  'DQY': 0.0,
  'DPY': 0.0,
  'TYPE': 'MONITOR',
  'RC': None},
 'V_BM_1': {'S': 9.0,
  'BX': 9.120629409314306,
  'AX': 1.146568740048649,
  'FX': 0.9936320883668346,
  'DQX': 1.992896322252739,
  'DPX': -0.2138526637243729,
  'BY': 10.91413760701323,
  'AY': -1.284847009200136,
  'FY': 1.001448183140224,
  'DQY': 0.0,
  'DPY': 0.0,
  'TYPE': 'VIRTUAL',
  'RC': None},
 'M_DR_3': {'S': 10.5,
  'BX': 5.980356479284523,
  'AX': 0.852404034821221,
  'FX': 1.197543640504269,
  'DQX': 1.744126672993481,
  'DPX': -0.1561981825607101,
  'BY': 15.31515988971355,
  'AY': -1.649167846239906,
  'FY': 1.11773084119158,
  'DQY': 0.0,
  'DPY': 0.0,
  'TYPE': 'MONITOR',
  'RC': None},
 'TAIL': {'S': 12.0,
  'BX': 4.287017734831201,
  'AX': 5.273559366969493e-16,
  'FX': 1.50436222394872,
  'DQX': 1.549040841795901,
  'DPX': -1.110223024625157e-16,
  'BY': 19.81848926815184,
  'AY': 1.387778780781446e-16,
  'FY': 1.202262333729486,
  'DQY': 0.0,
  'DPY': 0.0,
  'TYPE': 'VIRTUAL',
  'RC': None}}
[10]:
# RC parameter from lattice data can be added to model table
# Configuration table can be saved using util.save

table = add_rc(table, lattice)
table
[10]:
{'HEAD': {'S': 0.0,
  'BX': 4.287017734831204,
  'AX': -1.321304551729011e-16,
  'FX': 0.0,
  'DQX': 1.549040841795901,
  'DPX': 2.775557561562891e-17,
  'BY': 19.81848926815186,
  'AY': 6.545730027929358e-16,
  'FY': 0.0,
  'DQY': 0.0,
  'DPY': 0.0,
  'TYPE': 'VIRTUAL',
  'RC': None},
 'M': {'S': 0.0,
  'BX': 4.287017734831204,
  'AX': -1.321304551729011e-16,
  'FX': 0.0,
  'DQX': 1.549040841795901,
  'DPX': 2.775557561562891e-17,
  'BY': 19.81848926815186,
  'AY': 6.545730027929358e-16,
  'FY': 0.0,
  'DQY': 0.0,
  'DPY': 0.0,
  'TYPE': 'MONITOR',
  'RC': None},
 'M_DR': {'S': 1.5,
  'BX': 5.980356479284525,
  'AX': -0.8524040348212205,
  'FX': 0.3068185834444499,
  'DQX': 1.744126672993481,
  'DPX': 0.1561981825607101,
  'BY': 15.31515988971357,
  'AY': 1.649167846239909,
  'FY': 0.08453149253790615,
  'DQY': 0.0,
  'DPY': 0.0,
  'TYPE': 'MONITOR',
  'RC': ['DR', {'KIND': 'DRIF', 'L': 2}]},
 'V_BM': {'S': 3.0,
  'BX': 9.120629409314306,
  'AX': -1.146568740048648,
  'FX': 0.5107301355818847,
  'DQX': 1.99289632225274,
  'DPX': 0.2138526637243729,
  'BY': 10.91413760701325,
  'AY': 1.284847009200138,
  'FY': 0.2008141505892614,
  'DQY': 0.0,
  'DPY': 0.0,
  'TYPE': 'VIRTUAL',
  'RC': ['BM',
   {'KIND': 'CSBEND', 'ANGLE': 0.17453292519943295, 'FINT': 0, 'L': 1.0}]},
 'M_DR_1': {'S': 4.5,
  'BX': 12.70896979491113,
  'AX': -1.336396672789745,
  'FX': 0.6500917126795847,
  'DQX': 2.383785789407781,
  'DPX': 0.2698795969893243,
  'BY': 7.606077834992553,
  'AY': 0.9205261718281931,
  'FY': 0.3661997933967667,
  'DQY': 0.0,
  'DPY': 0.0,
  'TYPE': 'MONITOR',
  'RC': ['DR', {'KIND': 'DRIF', 'L': 2}]},
 'V_QF': {'S': 6.0,
  'BX': 16.39200719526896,
  'AX': 3.608224830031759e-16,
  'FX': 0.7521811119743597,
  'DQX': 2.721417822995524,
  'DPX': -5.551115123125783e-17,
  'BY': 5.674619589422572,
  'AY': 8.326672684688672e-16,
  'FY': 0.6011311668647427,
  'DQY': 0.0,
  'DPY': 0.0,
  'TYPE': 'VIRTUAL',
  'RC': ['QF', {'KIND': 'QUAD', 'K1': 0.2, 'L': 1.0}]},
 'M_DR_2': {'S': 7.5,
  'BX': 12.70896979491113,
  'AX': 1.336396672789746,
  'FX': 0.8542705112691347,
  'DQX': 2.383785789407781,
  'DPX': -0.2698795969893243,
  'BY': 7.606077834992548,
  'AY': -0.920526171828191,
  'FY': 0.8360625403327188,
  'DQY': 0.0,
  'DPY': 0.0,
  'TYPE': 'MONITOR',
  'RC': ['DR', {'KIND': 'DRIF', 'L': 2}]},
 'V_BM_1': {'S': 9.0,
  'BX': 9.120629409314306,
  'AX': 1.146568740048649,
  'FX': 0.9936320883668346,
  'DQX': 1.992896322252739,
  'DPX': -0.2138526637243729,
  'BY': 10.91413760701323,
  'AY': -1.284847009200136,
  'FY': 1.001448183140224,
  'DQY': 0.0,
  'DPY': 0.0,
  'TYPE': 'VIRTUAL',
  'RC': ['BM',
   {'KIND': 'CSBEND', 'ANGLE': 0.17453292519943295, 'FINT': 0, 'L': 1.0}]},
 'M_DR_3': {'S': 10.5,
  'BX': 5.980356479284523,
  'AX': 0.852404034821221,
  'FX': 1.197543640504269,
  'DQX': 1.744126672993481,
  'DPX': -0.1561981825607101,
  'BY': 15.31515988971355,
  'AY': -1.649167846239906,
  'FY': 1.11773084119158,
  'DQY': 0.0,
  'DPY': 0.0,
  'TYPE': 'MONITOR',
  'RC': ['DR', {'KIND': 'DRIF', 'L': 2}]},
 'TAIL': {'S': 12.0,
  'BX': 4.287017734831201,
  'AX': 5.273559366969493e-16,
  'FX': 1.50436222394872,
  'DQX': 1.549040841795901,
  'DPX': -1.110223024625157e-16,
  'BY': 19.81848926815184,
  'AY': 1.387778780781446e-16,
  'FY': 1.202262333729486,
  'DQY': 0.0,
  'DPY': 0.0,
  'TYPE': 'VIRTUAL',
  'RC': None}}

Example-04: Transformations benchmark (PTC)

[1]:
# In this example various symplectic transformations are compared with corresponding MADX-PTC transformations
[2]:
# calibration

import torch

from model.library.transformations import calibration_forward
from model.library.transformations import calibration_inverse

gxx = torch.tensor(1.005, dtype=torch.float64)
gxy = torch.tensor(0.001, dtype=torch.float64)
gyx = torch.tensor(0.005, dtype=torch.float64)
gyy = torch.tensor(0.955, dtype=torch.float64)

state = torch.tensor([0.01, -0.005, -0.01, 0.005], dtype=torch.float64)

res = calibration_forward(state, gxx, gxy, gyx, gyy)
res = calibration_inverse(res, gxx, gxy, gyx, gyy)

print(state.tolist())
print(res.tolist())
print((state - res).tolist())
[0.01, -0.005, -0.01, 0.005]
[0.010000000000000002, -0.005, -0.010000000000000002, 0.005]
[-1.734723475976807e-18, 0.0, 1.734723475976807e-18, 0.0]
[3]:
# drift

from pathlib import Path
from os import system

import torch
from model.library.transformations import drift

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

dp = torch.tensor(0.005, dtype=torch.float64)
length = torch.tensor(1.0, dtype=torch.float64)
state = torch.tensor([0.01, -0.005, -0.01, 0.005], dtype=torch.float64)

qx, px, qy, py = state.tolist()

code = f"""
mag:drift,l={length.item()} ;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map  ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact=false ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;
ptc_align;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;
ptc_track,icase=6,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}};
ptc_track_end;
ptc_end;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)
res = drift(state, dp, length)

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.005024875621890547, -0.005, -0.005024875621890547, 0.005]
[0.005024875621890547, -0.005, -0.005024875621890547, 0.005]
[0.0, 0.0, 0.0, 0.0]
[4]:
# drift (with kinematic)

from pathlib import Path
from os import system

import torch
from model.library.transformations import drift
from model.library.transformations import kinematic

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

dp = torch.tensor(0.005, dtype=torch.float64)
length = torch.tensor(1.0, dtype=torch.float64)
state = torch.tensor([0.01, -0.005, -0.01, 0.005], dtype=torch.float64)

qx, px, qy, py = state.tolist()

code = f"""
mag:drift,l={length.item()} ;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map  ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact=true ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;
ptc_align;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;
ptc_track,icase=6,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}};
ptc_track_end;
ptc_end;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)
state = drift(state, dp, length)
state = kinematic(state, dp, length)
res = state

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.005024752473723395, -0.005, -0.005024752473723395, 0.005]
[0.005024752473723394, -0.005, -0.005024752473723394, 0.005]
[8.673617379884035e-19, 0.0, -8.673617379884035e-19, 0.0]
[5]:
# corrector

from pathlib import Path
from os import system

import torch
from model.library.transformations import corrector

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

dp = torch.tensor(0.005, dtype=torch.float64)
kx = -torch.tensor(0.1, dtype=torch.float64)
ky = +torch.tensor(0.1, dtype=torch.float64)
state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)

qx, px, qy, py = state.tolist()

code = f"""
magx:hkicker,l=0.0,kick={kx.item()};
magy:vkicker,l=0.0,kick={ky.item()};
map:line=(magx, magy) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map  ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact=false ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;
ptc_align;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;
ptc_track,icase=6,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}};
ptc_track_end;
ptc_end;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)
res = corrector(state, kx, ky)

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.0, -0.1, 0.0, 0.1]
[0.0, -0.1, 0.0, 0.1]
[0.0, 0.0, 0.0, 0.0]
[6]:
# focusing quadrupole

from pathlib import Path
from os import system

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

import torch
from model.library.transformations import fquad

kn = torch.tensor(1.0, dtype=torch.float64)
dp = torch.tensor(0.005, dtype=torch.float64)
length = torch.tensor(1.0, dtype=torch.float64)
state = torch.tensor([0.01, -0.005, -0.01, 0.005], dtype=torch.float64)

qx, px, qy, py = state.tolist()

code = f"""
mag:quadrupole,l={length.item()}, k1={kn.item()};
map:line=(mag) ;
beam,energy=1.0E+9,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map  ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact=false ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;
ptc_align;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;
ptc_track,icase=6,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}};
ptc_track_end;
ptc_end;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)
res = fquad(state, kn.abs(), dp, length)

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.0012338134845463512, -0.011134185772061586, -0.009559363509475737, -0.004042070986490389]
[0.0012338134845463642, -0.01113418577206157, -0.009559363509475702, -0.0040420709864903495]
[-1.3010426069826053e-17, -1.5612511283791264e-17, -3.469446951953614e-17, -3.9898639947466563e-17]
[7]:
# defocusing quadrupole

from pathlib import Path
from os import system

import torch
from model.library.transformations import dquad

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

kn = torch.tensor(-1.0, dtype=torch.float64)
dp = torch.tensor(0.005, dtype=torch.float64)
length = torch.tensor(1.0, dtype=torch.float64)
state = torch.tensor([0.01, -0.005, -0.01, 0.005], dtype=torch.float64)

qx, px, qy, py = state.tolist()

code = f"""
mag:quadrupole,l={length.item()}, k1={kn.item()};
map:line=(mag) ;
beam,energy=1.0E+9,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map  ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact=false ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;
ptc_align;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;
ptc_track,icase=6,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}};
ptc_track_end;
ptc_end;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)
res = dquad(state, kn.abs(), dp, length)

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.009559363509475737, 0.004042070986490389, -0.0012338134845463512, 0.011134185772061586]
[0.009559363509475702, 0.0040420709864903495, -0.0012338134845463642, 0.01113418577206157]
[3.469446951953614e-17, 3.9898639947466563e-17, 1.3010426069826053e-17, 1.5612511283791264e-17]
[8]:
# generic quadrupole

from pathlib import Path
from os import system

import torch
from model.library.transformations import quadrupole

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

kn = torch.tensor(-1.0, dtype=torch.float64)
ks = torch.tensor(+1.0, dtype=torch.float64)
dp = torch.tensor(0.005, dtype=torch.float64)
length = torch.tensor(1.0, dtype=torch.float64)
state = torch.tensor([0.01, -0.005, -0.01, 0.005], dtype=torch.float64)

qx, px, qy, py = state.tolist()

code = f"""
mag: quadrupole, l={length.item()},k1={kn.item()},k1s={ks.item()};
map:line=(mag) ;
beam,energy=1.0E+9,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map  ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact=false ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;
ptc_align;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;
ptc_track,icase=6,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}};
ptc_track_end;
ptc_end;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)
res = quadrupole(state, kn, ks, dp, length)

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.00576876084434182, -0.0020884471766207963, 0.002582224887596925, 0.017416187766128223]
[0.005768760844341825, -0.0020884471766207725, 0.0025822248875969544, 0.0174161877661282]
[-5.204170427930421e-18, -2.3852447794681098e-17, -2.949029909160572e-17, 2.42861286636753e-17]
[9]:
# generic linear transformation

import torch
from model.library.transformations import quadrupole
from model.library.transformations import linear

kn = torch.tensor(-1.0, dtype=torch.float64)
ks = torch.tensor(+1.0, dtype=torch.float64)
dp = torch.tensor(0.005, dtype=torch.float64)
length = torch.tensor(1.0, dtype=torch.float64)
state = torch.tensor([0.01, -0.005, -0.01, 0.005], dtype=torch.float64)

vector = torch.zeros_like(state)
matrix = torch.func.jacrev(quadrupole)(0.0*state, kn, ks, dp, length)

ref = quadrupole(state, kn, ks, dp, length)
res = linear(state, vector, matrix)

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())
[0.005768760844341825, -0.0020884471766207725, 0.0025822248875969544, 0.0174161877661282]
[0.005768760844341825, -0.0020884471766207725, 0.0025822248875969544, 0.0174161877661282]
[0.0, 0.0, 0.0, 0.0]
[10]:
# quadrupole kick

from pathlib import Path
from os import system

import torch
from model.library.transformations import gradient

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

kn = torch.tensor(1.5, dtype=torch.float64)
ks = torch.tensor(1.0, dtype=torch.float64)
dp = torch.tensor(0.005, dtype=torch.float64)
length = torch.tensor(1.0, dtype=torch.float64)
state = torch.tensor([0.01, -0.005, -0.05, 0.001], dtype=torch.float64)

qx, px, qy, py = state.tolist()

code = f"""
mag: multipole,knl={{0.0,{(kn*length).item()}}},ksl={{0.0,{(ks*length).item()}}};
map:line=(mag) ;
beam,energy=1.0E+9,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map  ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact=false ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;
ptc_track,icase=6,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}};
ptc_track_end;
ptc_end;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)
res = gradient(state, kn, ks, length)

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.01, -0.07, -0.05, -0.06400000000000002]
[0.01, -0.07, -0.05, -0.06400000000000002]
[0.0, 0.0, 0.0, 0.0]
[11]:
# sextupole kick

from pathlib import Path
from os import system

import torch
from model.library.transformations import sextupole

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

ks = torch.tensor(0.5, dtype=torch.float64)
dp = torch.tensor(0.005, dtype=torch.float64)
length = torch.tensor(1.0, dtype=torch.float64)
state = torch.tensor([0.01, -0.005, -0.05, 0.001], dtype=torch.float64)

qx, px, qy, py = state.tolist()

code = f"""
mag: multipole,knl={{0.0,0.0,{(ks*length).item()}}};
map:line=(mag) ;
beam,energy=1.0E+9,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map  ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact=false ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;
ptc_track,icase=6,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}};
ptc_track_end;
ptc_end;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)
res = sextupole(state, ks, length)

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.01, -0.0044, -0.05, 0.00075]
[0.01, -0.0044, -0.05, 0.00075]
[0.0, 0.0, 0.0, 0.0]
[12]:
# octupole kick

import torch
from model.library.transformations import octupole

from pathlib import Path
from os import system

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

ko = torch.tensor(5.0, dtype=torch.float64)
dp = torch.tensor(0.005, dtype=torch.float64)
length = torch.tensor(1.0, dtype=torch.float64)
state = torch.tensor([0.01, -0.005, -0.05, 0.001], dtype=torch.float64)

qx, px, qy, py = state.tolist()

code = f"""
mag: multipole,knl={{0.0,0.0,0.0,{(ko*length).item()}}};
map:line=(mag) ;
beam,energy=1.0E+9,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map  ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact=false ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;
ptc_track,icase=6,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}};
ptc_track_end;
ptc_end;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)
res = octupole(state, ko, length)

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.01, -0.004938333333333334, -0.05, 0.0010916666666666668]
[0.01, -0.004938333333333334, -0.05, 0.0010916666666666668]
[0.0, 0.0, 0.0, 0.0]
[13]:
# pure dipole (no edge effects if e1 = e2 = 0)

from pathlib import Path
from os import system

import torch
from model.library.transformations import dipole

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

angle = torch.tensor(0.1, dtype=torch.float64 )
dp = torch.tensor(0.005, dtype=torch.float64)
length = torch.tensor(1.0, dtype=torch.float64)
state = torch.tensor([0.01, -0.005, -0.01, 0.005], dtype=torch.float64)

qx, px, qy, py = state.tolist()

code = f"""
mag: sbend, l={length.item()}, angle={angle.item()},k1=0.0,k1s=0.0,e1=0.0,e2=0.0,kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+9,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map  ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact=false ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;
ptc_align;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;
ptc_track,icase=6,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}};
ptc_track_end;
ptc_end;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)
res = dipole(state, length/angle, dp, length)

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.005231962156363519, -0.004575808017791941, -0.005024875621890682, 0.005]
[0.005231962156363559, -0.004575808017791928, -0.005024875621890547, 0.005]
[-3.9898639947466563e-17, -1.3010426069826053e-17, -1.3530843112619095e-16, 0.0]
[14]:
# combined function dipole (no edge effects if e1 = e2 = 0)

from pathlib import Path
from os import system

import torch
from model.library.transformations import bend

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

angle = torch.tensor(0.1, dtype=torch.float64)
kn = torch.tensor(-1.0, dtype=torch.float64 )
ks = torch.tensor(0.5, dtype=torch.float64 )
dp = torch.tensor(0.005, dtype=torch.float64)
length = torch.tensor(1.0, dtype=torch.float64)
state = torch.tensor([0.01, -0.005, -0.01, 0.005], dtype=torch.float64)

qx, px, qy, py = state.tolist()

code = f"""
mag: sbend, l={length.item()}, angle={angle.item()},k1={kn.item()},k1s={ks.item()},e1=0.0,e2=0.0,kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+9,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map  ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact=false ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;
ptc_align;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;
ptc_track,icase=6,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}};
ptc_track_end;
ptc_end;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)
res = bend(state, length/angle, kn, ks, dp, length)

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.007795176039169286, 0.0011097346020520168, 0.0007677684542948365, 0.01462661777023533]
[0.0077951760391692755, 0.0011097346020520064, 0.0007677684542948021, 0.014626617770235346]
[1.0408340855860843e-17, 1.0408340855860843e-17, 3.436920886779049e-17, -1.5612511283791264e-17]
[15]:
# combined function dipole with edge effects

from pathlib import Path
from os import system

import torch
from model.library.transformations import bend
from model.library.transformations import wedge

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

angle = torch.tensor(0.1, dtype=torch.float64)
kn = torch.tensor(-1.0, dtype=torch.float64)
ks = torch.tensor(0.5, dtype=torch.float64)
dp = torch.tensor(0.005, dtype=torch.float64)
e1 = torch.tensor(0.02, dtype=torch.float64)
e2 = torch.tensor(-0.03, dtype=torch.float64)
length = torch.tensor(2.5, dtype=torch.float64)
state = torch.tensor([0.01, -0.005, -0.01, 0.005], dtype=torch.float64)

qx, px, qy, py = state.tolist()

code = f"""
mag: sbend, l={length.item()}, angle={angle.item()},k1={kn.item()},k1s={ks.item()},e1={e1.item()},e2={e2.item()},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+9,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map  ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact=false ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=false ;
ptc_align;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;
ptc_track,icase=6,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}};
ptc_track_end;
ptc_end;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

state = wedge(state, e1, length/angle)
state = bend(state, length/angle, kn + 1.0E-16, ks + 1.0E-16, dp, length)
res = wedge(state, e2, length/angle)

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.02505664248885726, 0.02888471799554683, 0.01948347950480536, 0.007752605073815676]
[0.025056642488857035, 0.028884717995546608, 0.019483479504805307, 0.007752605073815616]
[2.255140518769849e-16, 2.220446049250313e-16, 5.204170427930421e-17, 5.984795992119984e-17]
[16]:
# translations (exact alignment, straight layout, act on a thin representation at the entrance frame)

from pathlib import Path
from os import system

import torch
from model.library.transformations import drift
from model.library.transformations import quadrupole
from model.library.transformations import tx, ty, tz

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

angle = torch.tensor(0.1, dtype=torch.float64)
kn = torch.tensor(-1.0, dtype=torch.float64)
ks = torch.tensor(0.5, dtype=torch.float64)
dp = torch.tensor(0.005, dtype=torch.float64)
length = torch.tensor(2.5, dtype=torch.float64)
state = torch.tensor([0.01, -0.005, -0.015, 0.0025], dtype=torch.float64)

dx = torch.tensor(0.01, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

qx, px, qy, py = state.tolist()

code = f"""
dr: drift, l=1.0;
mag: quadrupole, l={length.item()},k1={kn.item()},k1s={ks.item()};
map:line=(dr, mag, dr) ;
beam,energy=1.0E+9,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map  ;
select,flag=error,pattern="mag";
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()};
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact=false ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;
ptc_track,icase=6,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}};
ptc_track_end;
ptc_end;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

state = drift(state, dp, 1.0)
state = tx(state, +dx)
state = ty(state, +dy)
state = tz(state, +dz, dp)
state = quadrupole(state, kn, ks, dp, length)
state = tz(state, -dz, dp)
state = ty(state, -dy)
state = tx(state, -dx)
state = drift(state, dp, 1.0)
res = state

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[-0.08749190399027759, -0.05149961687573391, -0.05640828383591964, -0.019912883990595223]
[-0.08749190399027745, -0.05149961687573378, -0.05640828383591953, -0.01991288399059518]
[-1.3877787807814457e-16, -1.249000902703301e-16, -1.1102230246251565e-16, -4.163336342344337e-17]
[17]:
# translations + rotations (exact alignment, straight layout, act on a thin representation at the entrance frame)

from pathlib import Path
from os import system

import torch
from model.library.transformations import quadrupole
from model.library.transformations import tx, ty, tz
from model.library.transformations import rx, ry, rz

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

kn = torch.tensor(-1.0, dtype=torch.float64)
ks = torch.tensor(0.5, dtype=torch.float64)
dp = torch.tensor(0.005, dtype=torch.float64)

length = torch.tensor(2.5, dtype=torch.float64)
state = torch.tensor([0.01, -0.005, -0.015, 0.0025], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

qx, px, qy, py = state.tolist()

code = f"""
mag:quadrupole,l={length.item()},k1={kn.item()},k1s={ks.item()};
map:line=(mag) ;
beam,energy=1.0E+9,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map  ;
select,flag=error,pattern="mag";
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()};
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact=false ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true;
ptc_align;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;
ptc_track,icase=6,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}};
ptc_track_end;
ptc_end;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

state = tx(state, +dx)
state = ty(state, +dy)
state = tz(state, +dz, dp)

state = rx(state, +wx, dp)
state = ry(state, +wy, dp)
state = rz(state, +wz)

state = quadrupole(state, kn + 1.0E-16, ks + 1.0E-16, dp, length)

state = tz(state, -length, dp)

state = rz(state, -wz)
state = ry(state, -wy, dp)
state = rx(state, -wx, dp)

state = tz(state, -dz, dp)
state = ty(state, -dy)
state = tx(state, -dx)

state = tz(state, +length, dp)

res = state

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[-0.19014967849021078, -0.2611883141418841, -0.10367399923762718, -0.09101974096085057]
[-0.19014967849022751, -0.2611883141418846, -0.10367399923763249, -0.09101974096085078]
[1.6736612096224235e-14, 4.996003610813204e-16, 5.3013149425851225e-15, 2.0816681711721685e-16]
[18]:
# translations + rotations (exact alignment, curved layout, act on a thin representation at the entrance frame)

from pathlib import Path
from os import system

import torch
from model.library.transformations import bend
from model.library.transformations import wedge
from model.library.transformations import tx, ty, tz
from model.library.transformations import rx, ry, rz

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

angle = torch.tensor(0.1, dtype=torch.float64)
kn = torch.tensor(-1.0, dtype=torch.float64)
ks = torch.tensor(0.5, dtype=torch.float64)
dp = torch.tensor(0.005, dtype=torch.float64)
e1 = torch.tensor(0.005, dtype=torch.float64)
e2 = torch.tensor(0.005, dtype=torch.float64)

length = torch.tensor(2.5, dtype=torch.float64)
state = torch.tensor([0.01, -0.005, -0.015, 0.0025], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.01, dtype=torch.float64)

qx, px, qy, py = state.tolist()

code = f"""
mag:sbend,l={length.item()},angle={angle.item()},k1={kn.item()},k1s={ks.item()},e1={e1.item()},e2={e2.item()},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+9,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map  ;
select,flag=error,pattern="mag";
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()};
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact=false ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true;
ptc_align;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp.item()},t=0.0 ;
ptc_track,icase=6,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}};
ptc_track_end;
ptc_end;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

state = tx(state, +dx)
state = ty(state, +dy)
state = tz(state, +dz, dp)

state = rx(state, +wx, dp)
state = ry(state, +wy, dp)
state = rz(state, +wz)

state = wedge(state, e1, length/angle)
state = bend(state, length/angle, kn + 1.0E-16, ks + 1.0E-16, dp, length)
state = wedge(state, e2, length/angle)

state = ry(state, +angle/2, dp)
state = tz(state, -2.0*length/angle*(angle/2.0).sin(), dp)
state = ry(state, +angle/2, dp)

state = rz(state, -wz)
state = ry(state, -wy, dp)
state = rx(state, -wx, dp)

state = tz(state, -dz, dp)
state = ty(state, -dy)
state = tx(state, -dx)

state = ry(state, -angle/2, dp)
state = tz(state, +2.0*length/angle*(angle/2.0).sin(), dp)
state = ry(state, -angle/2, dp)

res = state

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[-0.20185873553633701, -0.27629369448112884, -0.08354566211146026, -0.06880534321809641]
[-0.20185873553633754, -0.2762936944811289, -0.08354566211146025, -0.06880534321809621]
[5.273559366969494e-16, 5.551115123125783e-17, -1.3877787807814457e-17, -1.942890293094024e-16]
[19]:
# exact sector bend (without fringe)

from pathlib import Path
from os import system

import torch
from model.library.transformations import sector_bend

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

length = 2.5
angle = 0.1
e1 = 0.0
e2 = 0.0
dp = 0.005

state = torch.tensor([0.01, -0.0005, -0.01, 0.0005], dtype=torch.float64)
qx, px, qy, py = state.tolist()

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},kill_ent_fringe=true,kill_exi_fringe=true;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map  ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact=true ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=6,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}};
ptc_track_end;
ptc_end;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

length = torch.tensor(length, dtype=torch.float64)
angle = torch.tensor(angle, dtype=torch.float64)
e1 = torch.tensor(e1, dtype=torch.float64)
e2 = torch.tensor(e2, dtype=torch.float64)
dp = torch.tensor(dp, dtype=torch.float64)

state = sector_bend(state, length/angle, dp, length)
res = state

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.00932966343120269, -3.829320024990772e-05, -0.008755742622854638, 0.0005]
[0.009329663431209667, -3.829320024994158e-05, -0.008755742622854576, 0.0005]
[-6.977057820378718e-15, 3.386098909943791e-17, -6.245004513516506e-17, 0.0]
[20]:
# exact sector bend (with fringe)

from pathlib import Path
from os import system

import torch
from model.library.transformations import sector_bend
from model.library.transformations import sector_bend_fringe

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

length = 2.5
angle = 0.1
e1 = 0.0
e2 = 0.0
dp = 0.005

state = torch.tensor([0.01, -0.0005, -0.01, 0.0005], dtype=torch.float64)
qx, px, qy, py = state.tolist()

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map  ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact=true ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=6,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}};
ptc_track_end;
ptc_end;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

length = torch.tensor(length, dtype=torch.float64)
angle = torch.tensor(angle, dtype=torch.float64)
e1 = torch.tensor(e1, dtype=torch.float64)
e2 = torch.tensor(e2, dtype=torch.float64)
dp = torch.tensor(dp, dtype=torch.float64)

state = sector_bend_fringe(state, +length/angle, dp)
state = sector_bend(state, length/angle, dp, length)
state = sector_bend_fringe(state, -length/angle, dp)
res = state

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.00933011773950498, -3.830113730822001e-05, -0.008756237750526722, 0.0004998143432367524]
[0.009330117739498277, -3.8301137308233764e-05, -0.008756237750526754, 0.0004998143432367524]
[6.702971511174383e-15, 1.3755815063409838e-17, 3.122502256758253e-17, 0.0]
[21]:
# exact sector bend (with fringe and wedges)

from pathlib import Path
from os import system

import torch
from model.library.transformations import sector_bend
from model.library.transformations import sector_bend_fringe
from model.library.transformations import sector_bend_wedge
from model.library.transformations import polar

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

length = 2.5
angle = 0.1
e1 = 0.01
e2 = -0.01
dp = 0.005

state = torch.tensor([0.01, -0.0005, -0.01, 0.0005], dtype=torch.float64)
qx, px, qy, py = state.tolist()

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map  ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact=true ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=6,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}};
ptc_track_end;
ptc_end;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

length = torch.tensor(length, dtype=torch.float64)
angle = torch.tensor(angle, dtype=torch.float64)
e1 = torch.tensor(e1, dtype=torch.float64)
e2 = torch.tensor(e2, dtype=torch.float64)
dp = torch.tensor(dp, dtype=torch.float64)

state = polar(state, e1, dp)
state = sector_bend_fringe(state, +length/angle, dp)
state = sector_bend_wedge(state, -e1, length/angle, dp)
state = sector_bend(state, length/angle, dp, length)
state = sector_bend_wedge(state, -e2, length/angle, dp)
state = sector_bend_fringe(state, -length/angle, dp)
state = polar(state, e2, dp)
res = state

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.009340060895391944, -3.805697197357787e-05, -0.00874628326389246, 0.0005003157283896672]
[0.009340060895401722, -3.805697197360909e-05, -0.00874628326389242, 0.0005003157283896674]
[-9.778636234081262e-15, 3.1218246304004493e-17, -3.9898639947466563e-17, -1.0842021724855044e-19]
[22]:
# linear wedge matrix

from model.library.transformations import wedge
from model.library.transformations import polar
from model.library.transformations import sector_bend_fringe
from model.library.transformations import sector_bend_wedge

length = 2.5
angle = 0.1
e1 = 0.01
e2 = -0.01
dp = 0.005

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)

length = torch.tensor(length, dtype=torch.float64)
angle = torch.tensor(angle, dtype=torch.float64)
e1 = torch.tensor(e1, dtype=torch.float64)
e2 = torch.tensor(e2, dtype=torch.float64)
dp = torch.tensor(dp, dtype=torch.float64)

def wedge_entrance(state):
    state = polar(state, e1, dp)
    state = sector_bend_fringe(state, +length/angle, dp)
    state = sector_bend_wedge(state, -e1, length/angle, dp)
    return state

def wedge_exit(state):
    state = sector_bend_wedge(state, -e2, length/angle, dp)
    state = sector_bend_fringe(state, -length/angle, dp)
    state = polar(state, e2, dp)
    return state

print(torch.func.jacrev(wedge)(state, e1, length/angle))
print(torch.func.jacrev(wedge_entrance)(state))
print()

print(torch.func.jacrev(wedge)(state, e2, length/angle))
print(torch.func.jacrev(wedge_exit)(state))
print()
tensor([[ 1.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],
        [ 4.0001e-04,  1.0000e+00,  0.0000e+00,  0.0000e+00],
        [ 0.0000e+00,  0.0000e+00,  1.0000e+00,  0.0000e+00],
        [ 0.0000e+00,  0.0000e+00, -4.0001e-04,  1.0000e+00]],
       dtype=torch.float64)
tensor([[ 1.0000e+00,  5.5508e-17,  0.0000e+00,  0.0000e+00],
        [ 4.0001e-04,  1.0000e+00,  0.0000e+00,  0.0000e+00],
        [ 0.0000e+00,  0.0000e+00,  1.0000e+00,  0.0000e+00],
        [ 0.0000e+00,  0.0000e+00, -4.0001e-04,  1.0000e+00]],
       dtype=torch.float64)

tensor([[ 1.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],
        [-4.0001e-04,  1.0000e+00,  0.0000e+00,  0.0000e+00],
        [ 0.0000e+00,  0.0000e+00,  1.0000e+00,  0.0000e+00],
        [ 0.0000e+00,  0.0000e+00,  4.0001e-04,  1.0000e+00]],
       dtype=torch.float64)
tensor([[ 1.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00],
        [-4.0001e-04,  1.0000e+00,  0.0000e+00,  0.0000e+00],
        [ 0.0000e+00,  0.0000e+00,  1.0000e+00,  5.5234e-17],
        [ 0.0000e+00,  0.0000e+00,  4.0001e-04,  1.0000e+00]],
       dtype=torch.float64)

Example-05: Drift (element)

[1]:
# Comparison of drift element with MADX-PTC and other features
[2]:
from pathlib import Path
from os import system

import torch
from model.library.drift import Drift
[3]:
# Tracking (paraxial)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = False
align = False

dp = 0.005
length = 1.5
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:drift,l={length} ;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Drift('D', length=length, dp=dp, exact=exact)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.0025373134328358204, -0.005, -0.003507462686567164, 0.001]
[0.0025373134328358204, -0.005, -0.0035074626865671645, 0.001]
[0.0, 0.0, 4.336808689942018e-19, 0.0]
[4]:
# Tracking (exact)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = True
align = False

dp = 0.005
length = 1.5
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:drift,l={length} ;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Drift('D', length=length, dp=dp, exact=exact)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.002537217378977325, -0.005, -0.003507443475795465, 0.001]
[0.002537217378977325, -0.005, -0.0035074434757954654, 0.001]
[0.0, 0.0, 4.336808689942018e-19, 0.0]
[5]:
# Tracking (exact, alignment)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = False
align = True

dp = 0.005
length = 1.5
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:drift,l={length} ;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Drift('D', length=length, dp=dp, exact=exact)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()

# Note, for some reason drift is not invariant under WX ans WY rotations in MADX
[0.0025372170792497145, -0.0049999999999999975, -0.003507395297032926, 0.0010000000000000009]
[0.0025372170792497166, -0.004999999999999999, -0.0035073952970329247, 0.0010000000000000018]
[-2.168404344971009e-18, 1.734723475976807e-18, -1.3010426069826053e-18, -8.673617379884035e-19]
[6]:
# Deviation/error variables

dp = 0.005
length = 1.5
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

D = Drift('D', length, dp)

# Each element has two variant of a call method
# In the first case only state is passed, it is transformed using parameters specified on initializaton

print(D(state))
print()

# Deviation errors can be also passed to call method
# These variables are added to corresponding parameters specified on initializaton
# For example, element lenght can changed

print(D(state, data={**D.data(), **{'dl': torch.tensor(-length, dtype=D.dtype)}}))
print()

# In the above D.data() creates default deviation dictionary (with zero values for each deviaton)
# {**D.data(), **{'dl': torch.tensor(-length, dtype=DR.dtype)}} replaces the 'dl' key value

# Additionaly, alignment errors are passed as deivation variables
# They are used if alignment flag is raised

print(D(state, data={**D.data(), **error}, alignment=True))
print()

# The following elements can be made equivalent using deviation variables

DA = Drift('DA', length, dp)
DB = Drift('DB', length - 0.1, dp)

print(DA(state) - DB(state, data={**DB.data(), **{'dl': torch.tensor(+0.1, dtype=DB.dtype)}}))
tensor([ 0.0025, -0.0050, -0.0035,  0.0010], dtype=torch.float64)

tensor([ 0.0100, -0.0050, -0.0050,  0.0010], dtype=torch.float64)

tensor([ 0.0025, -0.0050, -0.0035,  0.0010], dtype=torch.float64)

tensor([0., 0., 0., 0.], dtype=torch.float64)
[7]:
# Insertion element

# In this mode elements are treated as thin insertions (at the center)
# Using parameters specified on initialization, transport two matrices are computed
# These matrices are used to insert the element
# Input state is transformed from the element center to its entrance
# Next, transformation from the entrance frame to the exit frame is performed
# This transformation can contain errors
# The final step is to transform state from the exit frame back to the element center
# Without errors, this results in identity transformation for linear elements

dp = 0.0
length = 1.5
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

D = Drift('D', length, dp, exact=False, insertion=True)

# Identity transformation without errors

print(D(state) - state)

# Represents effect of an error

print(D(state, data={**D.data(), **{'dl': 0.1}}) - state)

# Exact tracking corresponds to inclusion of kinematic term as errors

D = Drift('D', length, dp, exact=True, insertion=True)

print(D(state) - state)
tensor([0., 0., 0., 0.], dtype=torch.float64)
tensor([-0.0005,  0.0000,  0.0001,  0.0000], dtype=torch.float64)
tensor([-9.7502e-08,  0.0000e+00,  1.9500e-08,  0.0000e+00],
       dtype=torch.float64)
[8]:
# Mapping over a set of initial conditions

# Call method can be used to map over a set of initial conditions
# Note, device can be set to cpu or gpu via base element classvariables

dp = 0.0
length = 1.5

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

D = Drift('D', length, dp, exact=True)

state = 1.0E-3*torch.randn((512, 4), dtype=D.dtype, device=D.device)

print(torch.vmap(D)(state).shape)

# To map over deviations parameters a wrapper function (or a lambda expression) can be used

def wrapper(state, dp):
    return D(state, data={**D.data(), **{'dp': dp}})

dp = 1.0E-3*torch.randn(512, dtype=D.dtype, device=D.device)

print(torch.vmap(wrapper)(state, dp).shape)
torch.Size([512, 4])
torch.Size([512, 4])
[9]:
# Differentiability

# Both call methods are differentiable
# Derivative with respect to state can be computed directly
# For deviation variables, wrapping is required

dp = 0.0
length = 1.5
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

D = Drift('D', length, dp, exact=False)

# Compute derivative with respect to state

print(torch.func.jacrev(D)(state))
print()

# Compute derivative with respect to a deviation variable

dl = torch.tensor(0.0, dtype=torch.float64)

def wrapper(state, dl):
    return D(state, data={**D.data(), **{'dl': dl}})

print(torch.func.jacrev(wrapper, 1)(state, dl))
print()

# Compositional derivative (compute derivative of jacobian trace with respect momentum deviation)

dp = torch.tensor(0.0, dtype=torch.float64)

def trace(state, dp):
    return (torch.func.jacrev(lambda state: D(state, data={**D.data(), **{'dp': dp}}))(state)).trace()

torch.func.jacrev(trace, 1)(state, dp)
tensor([[1.0000, 1.5000, 0.0000, 0.0000],
        [0.0000, 1.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 1.0000, 1.5000],
        [0.0000, 0.0000, 0.0000, 1.0000]], dtype=torch.float64)

tensor([-0.0050,  0.0000,  0.0010,  0.0000], dtype=torch.float64)

[9]:
tensor(-0., dtype=torch.float64)
[10]:
# Output at each step

# It is possible to collect output of state or tangent matrix at each integration step
# Number of integratin steps is controlled by ns parameter on initialization
# Alternatively, desired integration step length can be passed
# Number of integration steps is computed as ceil(length/ds)

dp = 0.0
length = 1.5
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

D = Drift('D', length, dp, exact=False, ns=10, output=True, matrix=True)

# Final state is still returned

print(D(state))

# Data is added to special attributes (state and tangent matrix)

print(D.container_output.shape)
print(D.container_matrix.shape)

# Number of integration steps can be changed

D.ns = 100

D(state)
print(D.container_output.shape)
print(D.container_matrix.shape)
tensor([ 0.0025, -0.0050, -0.0035,  0.0010], dtype=torch.float64)
torch.Size([10, 4])
torch.Size([10, 4, 4])
torch.Size([100, 4])
torch.Size([100, 4, 4])
[11]:
# Integration order is set on initialization (default value is zero)
# This order is related to difference order as 2n + 2
# Thus, zero corresponds to second order difference method

dp = 0.0
length = 1.5
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

D = Drift('D', length, dp, order=1, exact=True)

# For drift integration is performed only with exact flag
# In this case, kinematic term error is added
# This term actually commutes with paraxial drift map
# But integration is still performed for consistency with matrix-kick-matrix split
# Only one integration step is required to get exact result

D.ns = 1
ref = D(state)

D.ns = 10
res = D(state)

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())
print()

# Integrator parameters are stored in data attribute (if integration is actually performed)

maps, weights = D._data
print(maps)
print(weights)
[0.002499902498098708, -0.005, -0.003499980499619743, 0.001]
[0.002499902498098709, -0.005, -0.0034999804996197477, 0.001]
[-8.673617379884035e-19, 0.0, 4.7704895589362195e-18, 0.0]

[0, 1, 0, 1, 0, 1, 0]
[0.6756035959798289, 1.3512071919596578, -0.17560359597982877, -1.7024143839193153, -0.17560359597982877, 1.3512071919596578, 0.6756035959798289]

Example-06: Quadrupole (element)

[1]:
# Comparison of quadrupole element with MADX-PTC and other features
[2]:
from pathlib import Path
from os import system

import torch
from model.library.drift import Drift
from model.library.quadrupole import Quadrupole
[3]:
# Tracking (paraxial)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = False
align = False

kn = - 2.0
ks = + 1.5
dp = 0.005
length = 1.0
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:quadrupole,l={length},k1={kn},k1s={ks} ;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

Q = Quadrupole('Q', length=length, kn=kn, ks=ks, dp=dp, exact=exact)
res = Q(state, alignment=align, data={**Q.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.012268608165994052, 0.012991610983278109, 0.005825218798687177, 0.01752224400608683]
[0.012268608165994056, 0.012991610983278081, 0.005825218798687173, 0.017522244006086804]
[-3.469446951953614e-18, 2.7755575615628914e-17, 4.336808689942018e-18, 2.42861286636753e-17]
[4]:
# Tracking (exact)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = True
align = False

kn = - 2.0
ks = + 1.5
dp = 0.005
length = 1.0
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:quadrupole,l={length},k1={kn},k1s={ks} ;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

Q = Quadrupole('Q', length=length, kn=kn, ks=ks, dp=dp, exact=exact, order=5, ns=5)
res = Q(state, alignment=align, data={**Q.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.012269208914523159, 0.012992157908766264, 0.005826335256074007, 0.017521822791072554]
[0.012269208914522952, 0.012992157908766015, 0.005826335256074073, 0.017521822791072776]
[2.0643209364124004e-16, 2.498001805406602e-16, -6.591949208711867e-17, -2.220446049250313e-16]
[5]:
# Tracking (exact, alignment)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = True
align = True

kn = - 2.0
ks = + 1.5
dp = 0.005
length = 1.0
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:quadrupole,l={length},k1={kn},k1s={ks} ;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

Q = Quadrupole('Q', length=length, kn=kn, ks=ks, dp=dp, exact=exact, order=5, ns=5)
res = Q(state, alignment=align, data={**Q.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[-0.022075488016794924, -0.09165584224601611, -0.04570124622656498, -0.08629975808408008]
[-0.02207548801679271, -0.09165584224601468, -0.04570124622656417, -0.08629975808408101]
[-2.213507155346406e-15, -1.429412144204889e-15, -8.118505867571457e-16, 9.298117831235686e-16]
[6]:
# Deviation/error variables

kn = - 2.0
ks = + 1.5
dp = 0.005
length = 1.5
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

Q = Quadrupole('Q', length, kn, ks, dp)

# Each element has two variant of a call method
# In the first case only state is passed, it is transformed using parameters specified on initializaton

print(Q(state))
print()


# Deviation errors can be also passed to call method
# These variables are added to corresponding parameters specified on initializaton
# For example, element lenght can changed

print(Q(state, data={**Q.data(), **{'dl': -Q.length}}))
print()

# In the above Q.data() creates default deviation dictionary (with zero values for each deviaton)
# {**Q.data(), **{'dl': -Q.length}} replaces the 'dl' key value

# Additionaly, alignment errors are passed as deivation variables
# They are used if alignment flag is raised

print(Q(state, data={**Q.data(), **error}, alignment=True))
print()

# The following elements can be made equivalent using deviation variables

QA = Quadrupole('QA', length, kn, ks, dp)
QB = Quadrupole('QB', length - 0.1, kn, ks, dp)

print(QA(state) - QB(state, data={**QB.data(), **{'dl': torch.tensor(+0.1, dtype=QB.dtype)}}))

# Note, while in some cases float values can be passed as values to deviation variables
# The correct behaviour in guaranteed only for tensors
tensor([0.0242, 0.0380, 0.0152, 0.0200], dtype=torch.float64)

tensor([ 0.0100, -0.0050, -0.0050,  0.0010], dtype=torch.float64)

tensor([-0.0908, -0.2335, -0.0963, -0.1316], dtype=torch.float64)

tensor([0., 0., 0., 0.], dtype=torch.float64)
[7]:
# Insertion element

# In this mode elements are treated as thin insertions (at the center)
# Using parameters specified on initialization, transport two matrices are computed
# These matrices are used to insert the element
# Input state is transformed from the element center to its entrance
# Next, transformation from the entrance frame to the exit frame is performed
# This transformation can contain errors
# The final step is to transform state from the exit frame back to the element center
# Without errors, this results in identity transformation for linear elements

kn = - 2.0
ks = + 1.5
dp = 0.005
length = 1.5
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

Q = Quadrupole('Q', length, kn, ks, dp, exact=False, insertion=True)

# Identity transformation without errors

print(Q(state) - state)

# Represents effect of an error

print(Q(state, data={**Q.data(), **{'dl': 0.1, 'kn': -0.1}}) - state)

# Exact tracking corresponds to inclusion of kinematic term as errors

Q = Quadrupole('Q', length, kn, ks, dp, exact=True, insertion=True, ns=100, order=1)

print(Q(state) - state)
tensor([-5.2042e-18,  1.0408e-17, -3.4694e-18, -3.4694e-18],
       dtype=torch.float64)
tensor([-0.0002,  0.0037,  0.0003,  0.0031], dtype=torch.float64)
tensor([-2.2924e-06, -3.9787e-06, -9.4215e-07,  2.1943e-07],
       dtype=torch.float64)
[8]:
# Mapping over a set of initial conditions

# Call method can be used to map over a set of initial conditions
# Note, device can be set to cpu or gpu via base element classvariables

kn = - 2.0
ks = + 1.5
dp = 0.0
length = 1.5

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

Q = Quadrupole('Q', length, kn, ks, dp, exact=True)

state = 1.0E-3*torch.randn((512, 4), dtype=Q.dtype, device=Q.device)

print(torch.vmap(Q)(state).shape)

# To map over deviations parameters a wrapper function (or a lambda expression) can be used

def wrapper(state, dp):
    return Q(state, data={**Q.data(), **{'dp': dp}})

dp = 1.0E-3*torch.randn(512, dtype=Q.dtype, device=Q.device)

print(torch.vmap(wrapper)(state, dp).shape)
torch.Size([512, 4])
torch.Size([512, 4])
[9]:
# Differentiability

# Both call methods are differentiable
# Derivative with respect to state can be computed directly
# For deviation variables, wrapping is required

kn = - 2.0
ks = + 1.5
dp = 0.0
length = 1.5
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

Q = Quadrupole('Q', length, kn, ks, dp, exact=False)

# Compute derivative with respect to state

print(torch.func.jacrev(Q)(state))
print()

# Compute derivative with respect to a deviation variable

kn = torch.tensor(0.0, dtype=torch.float64)

def wrapper(state, kn):
    return Q(state, data={**Q.data(), **{'kn': kn}})

print(torch.func.jacrev(wrapper, 1)(state, kn))
print()

# Compositional derivative (compute derivative of jacobian trace with respect quadrupole strength)

length = 0.5
knf = +0.2
knd = -0.2

QF = Quadrupole('QF', length, knf)
QD = Quadrupole('QD', length, knd)
DR = Drift('DR', 5.0)

dknf = torch.tensor(0.0, dtype=torch.float64)
dknd = torch.tensor(0.0, dtype=torch.float64)
dkn = torch.stack([dknf, dknd])

def fodo(state, dkn):
    dknf, dknd = dkn
    state = QF(state, data={**QF.data(), **{'kn': dknf}})
    state = DR(state)
    state = QD(state, data={**QD.data(), **{'kn': dknd}})
    state = QD(state, data={**QD.data(), **{'kn': dknd}})
    state = DR(state)
    state = QF(state, data={**QF.data(), **{'kn': dknf}})
    return state

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)

def trace(dkn):
    return (torch.func.jacrev(fodo)(state, dkn)).trace()

torch.func.jacrev(trace)(dkn)
tensor([[ 4.7923,  3.0672,  1.8367,  0.8757],
        [ 7.4479,  4.7923,  2.8495,  1.8367],
        [ 1.8367,  0.8757, -0.1057,  0.7321],
        [ 2.8495,  1.8367, -0.1507, -0.1057]], dtype=torch.float64)

tensor([-0.0175, -0.0354, -0.0029, -0.0029], dtype=torch.float64)

[9]:
tensor([-12.7901,  12.7901], dtype=torch.float64)
[10]:
# Output at each step

# It is possible to collect output of state or tangent matrix at each integration step
# Number of integratin steps is controlled by ns parameter on initialization
# Alternatively, desired integration step length can be passed
# Number of integration steps is computed as ceil(length/ds)

kn = - 2.0
ks = + 1.5
dp = 0.0
length = 1.5
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

Q = Quadrupole('Q', length, kn, ks, dp, exact=False, ns=10, output=True, matrix=True)

# Final state is still returned

print(Q(state))

# Data is added to special attributes (state and tangent matrix)

print(Q.container_output.shape)
print(Q.container_matrix.shape)

# Number of integration steps can be changed

Q.ns = 100

Q(state)
print(Q.container_output.shape)
print(Q.container_matrix.shape)
tensor([0.0243, 0.0381, 0.0153, 0.0200], dtype=torch.float64)
torch.Size([10, 4])
torch.Size([10, 4, 4])
torch.Size([100, 4])
torch.Size([100, 4, 4])
[11]:
# Integration order is set on initialization (default value is zero)
# This order is related to difference order as 2n + 2
# Thus, zero corresponds to second order difference method

kn = - 2.0
ks = + 1.5
dp = 0.0
length = 1.5
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

Q = Quadrupole('Q', length, kn, ks, dp, order=1, exact=True)

# For quadrupole integration is performed only with exact flag
# In this case, kinematic term error is added

Q.ns = 1
ref = Q(state)

Q.ns = 10
res = Q(state)

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())
print()

# Integrator parameters are stored in data attribute (if integration is actually performed)

maps, weights = Q._data
print(maps)
print(weights)
[0.024284271092022615, 0.03811354319280181, 0.015254854998694, 0.01995832080452666]
[0.024286764143785347, 0.038112599406681526, 0.015255420211007705, 0.01995809327180016]
[-2.4930517627322346e-06, 9.437861202832298e-07, -5.652123137040582e-07, 2.2753272650027911e-07]

[0, 1, 0, 1, 0, 1, 0]
[0.6756035959798289, 1.3512071919596578, -0.17560359597982877, -1.7024143839193153, -0.17560359597982877, 1.3512071919596578, 0.6756035959798289]
[12]:
# Derivatives of twiss parameters

# pip install git+https://github.com/i-a-morozov/twiss.git@main

from twiss import twiss

length = 0.5
knf = +0.21
knd = -0.19

QF = Quadrupole('QF', length, knf)
QD = Quadrupole('QD', length, knd)
DR = Drift('DR', 5.0)

dknf = torch.tensor(0.0, dtype=torch.float64)
dknd = torch.tensor(0.0, dtype=torch.float64)
dkn = torch.stack([dknf, dknd])

def fodo(state, dkn):
    dknf, dknd = dkn
    state = QF(state, data={**QF.data(), **{'kn': dknf}})
    state = DR(state)
    state = QD(state, data={**QD.data(), **{'kn': dknd}})
    state = QD(state, data={**QD.data(), **{'kn': dknd}})
    state = DR(state)
    state = QF(state, data={**QF.data(), **{'kn': dknf}})
    return state

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)

def tune(dkn):
    matrix = torch.func.jacrev(fodo)(state, dkn)
    tune, *_ = twiss(matrix)
    return tune

# Compute tunes and jacobian

values = tune(dkn)
jacobian = torch.func.jacrev(tune)(dkn)

# Test jacobiant

print(values)
print(tune(dkn + 1.0E-3))
print(values + jacobian @ (dkn + 1.0E-3))
tensor([0.2107, 0.1703], dtype=torch.float64)
tensor([0.2126, 0.1681], dtype=torch.float64)
tensor([0.2126, 0.1681], dtype=torch.float64)

Example-07: Sextupole (element)

[1]:
# Comparison of sextupole element with MADX-PTC and other features
[2]:
from pathlib import Path
from os import system

import torch
from model.library.drift import Drift
from model.library.quadrupole import Quadrupole
from model.library.sextupole import Sextupole
[3]:
# Tracking (paraxial)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = False
align = False

ms = 10.0
dp = 0.005
length = 0.25
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag: sextupole, l={length},k2={ms} ;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

S = Sextupole('S', length=length, ms=ms, dp=dp, exact=exact, order=5, ns=10)
res = S(state, alignment=align, data={**S.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.008745689261382875, -0.005080234821325765, -0.00476590910682766, 0.0008855562050471031]
[0.008745689261382871, -0.00508023482132578, -0.00476590910682768, 0.0008855562050471008]
[3.469446951953614e-18, 1.5612511283791264e-17, 1.9949319973733282e-17, 2.2768245622195593e-18]
[4]:
# Tracking (exact)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = True
align = False

ms = 10.0
dp = 0.005
length = 0.25
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag: sextupole, l={length},k2={ms} ;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

S = Sextupole('S', length=length, ms=ms, dp=dp, exact=exact, order=5, ns=10)
res = S(state, alignment=align, data={**S.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.008745672936611987, -0.005080234654232574, -0.004765906047078759, 0.000885556338827788]
[0.00874567293661202, -0.0050802346542325495, -0.0047659060470788064, 0.0008855563388277869]
[-3.2959746043559335e-17, -2.42861286636753e-17, 4.7704895589362195e-17, 1.0842021724855044e-18]
[5]:
# Tracking (exact, alignment)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = True
align = True

ms = 10.0
dp = 0.005
length = 0.25
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag: sextupole, l={length},k2={ms} ;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

S = Sextupole('S', length=length, ms=ms, dp=dp, exact=exact, order=5, ns=10)
res = S(state, alignment=align, data={**S.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.008663885569968804, -0.006258052120536049, -0.004896687680053297, -0.000915022709372755]
[0.008663885569968922, -0.006258052120536033, -0.004896687680053253, -0.000915022709372733]
[-1.1796119636642288e-16, -1.6479873021779667e-17, -4.423544863740858e-17, -2.200930410145574e-17]
[6]:
# Deviation/error variables

ms = 10.0
dp = 0.005
length = 0.25
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

S = Sextupole('S', length, ms, dp)

# Each element has two variant of a call method
# In the first case only state is passed, it is transformed using parameters specified on initializaton

print(S(state))
print()

# Deviation errors can be also passed to call method
# These variables are added to corresponding parameters specified on initializaton
# For example, element lenght can changed

print(S(state, data={**S.data(), **{'dl': -S.length}}))
print()

# In the above S.data() creates default deviation dictionary (with zero values for each deviaton)
# {**S.data(), **{'dl': -S.length}} replaces the 'dl' key value

# Additionaly, alignment errors are passed as deivation variables
# They are used if alignment flag is raised

print(S(state, data={**S.data(), **error}, alignment=True))
print()

# The following elements can be made equivalent using deviation variables

SA = Sextupole('SA', length, ms, dp)
SB = Sextupole('SB', length - 0.1, ms, dp)

print(SA(state) - SB(state, data={**SB.data(), **{'dl': torch.tensor(+0.1, dtype=SB.dtype)}}))

# Note, while in some cases float values can be passed as values to deviation variables
# The correct behaviour in guaranteed only for tensors
tensor([ 0.0087, -0.0051, -0.0048,  0.0009], dtype=torch.float64)

tensor([ 0.0100, -0.0050, -0.0050,  0.0010], dtype=torch.float64)

tensor([ 0.0087, -0.0062, -0.0049, -0.0009], dtype=torch.float64)

tensor([0., 0., 0., 0.], dtype=torch.float64)
[7]:
# Insertion element

# In this mode elements are treated as thin insertions (at the center)
# Using parameters specified on initialization, transport two matrices are computed
# These matrices are used to insert the element
# Input state is transformed from the element center to its entrance
# Next, transformation from the entrance frame to the exit frame is performed
# This transformation can contain errors
# The final step is to transform state from the exit frame back to the element center
# Without errors, this results in identity transformation for linear elements

ms = 10.0
dp = 0.005
length = 1.5
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

S = Sextupole('S', length, ms, dp, exact=False, insertion=True)

# Since sextupole is a nonlinear element, insertion is an identity transformation only for zero strenght

print(S(state) - state)
print(S(state, data={**S.data(), **{'ms': -ms}}) - state)

# Represents effect of an error (any nonzero value of strengh or a change in other parameter)

print(S(state, data={**S.data(), **{'dl': 0.1}}) - state)

# Exact tracking corresponds to inclusion of kinematic term as errors

S = Sextupole('S', length, ms, dp, exact=True, insertion=True, ns=100, order=1)

print(S(state) - state)
tensor([ 0.0000, -0.0006,  0.0000, -0.0008], dtype=torch.float64)
tensor([0., 0., 0., 0.], dtype=torch.float64)
tensor([-5.2560e-04, -5.6465e-04,  6.1078e-05, -7.7234e-04],
       dtype=torch.float64)
tensor([-1.3875e-04, -5.5306e-04, -9.3764e-05, -7.7778e-04],
       dtype=torch.float64)
[8]:
# Mapping over a set of initial conditions

# Call method can be used to map over a set of initial conditions
# Note, device can be set to cpu or gpu via base element classvariables

ms = 10.0
dp = 0.0
length = 1.5

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

S = Sextupole('S', length, ms, dp, exact=True)

state = 1.0E-3*torch.randn((512, 4), dtype=S.dtype, device=S.device)

print(torch.vmap(S)(state).shape)

# To map over deviations parameters a wrapper function (or a lambda expression) can be used

def wrapper(state, dp):
    return S(state, data={**S.data(), **{'dp': dp}})

dp = 1.0E-3*torch.randn(512, dtype=S.dtype, device=S.device)

print(torch.vmap(wrapper)(state, dp).shape)
torch.Size([512, 4])
torch.Size([512, 4])
[9]:
# Differentiability

# Both call methods are differentiable
# Derivative with respect to state can be computed directly
# For deviation variables, wrapping is required

ms = 10.0
dp = 0.0
length = 1.5
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

S = Sextupole('S', length, ms, dp, exact=False)

# Compute derivative with respect to state

print(torch.func.jacrev(S)(state))
print()

# Compute derivative with respect to a deviation variable

ms = torch.tensor(0.0, dtype=torch.float64)

def wrapper(state, ms):
    return S(state, data={**S.data(), **{'ms': ms}})

print(torch.func.jacrev(wrapper, 1)(state, ms))
print()
tensor([[ 0.9297,  1.4473, -0.0478, -0.0359],
        [-0.0938,  0.9297, -0.0638, -0.0478],
        [-0.0478, -0.0359,  1.0703,  1.5527],
        [-0.0638, -0.0478,  0.0938,  1.0703]], dtype=torch.float64)

tensor([-1.1813e-05, -1.5750e-05, -2.9883e-05, -3.9844e-05],
       dtype=torch.float64)

[10]:
# Output at each step

# It is possible to collect output of state or tangent matrix at each integration step
# Number of integratin steps is controlled by ns parameter on initialization
# Alternatively, desired integration step length can be passed
# Number of integration steps is computed as ceil(length/ds)

ms = 10.0
dp = 0.0
length = 1.5
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

S = Sextupole('S', length, ms, dp, exact=False, ns=10, output=True, matrix=True)

# Final state is still returned

print(S(state))

# Data is added to special attributes (state and tangent matrix)

print(S.container_output.shape)
print(S.container_matrix.shape)

# Number of integration steps can be changed

S.ns = 100

S(state)
print(S.container_output.shape)
print(S.container_matrix.shape)
tensor([ 0.0023, -0.0052, -0.0039,  0.0006], dtype=torch.float64)
torch.Size([10, 4])
torch.Size([10, 4, 4])
torch.Size([100, 4])
torch.Size([100, 4, 4])
[11]:
# Integration order is set on initialization (default value is zero)
# This order is related to difference order as 2n + 2
# Thus, zero corresponds to second order difference method

ms = 10.0
dp = 0.0
length = 1.5
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

S = Sextupole('S', length, ms, dp, order=0, exact=True)

# For sextupole integration is always performed
# In exact case, kinematic term error is added

S.ns = 10
ref = S(state)

S.ns = 100
res = S(state)

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())
print()

# Integrator parameters are stored in data attribute (if integration is actually performed)

maps, weights = S._data
print(maps)
print(weights)
[0.0022880403040043407, -0.005176627884407675, -0.0038882248469424585, 0.0005834231060909752]
[0.0022871610802911667, -0.005176794799148338, -0.0038891562083014203, 0.0005832353574511327]
[8.792237131739246e-07, 1.6691474066278522e-07, 9.313613589618727e-07, 1.877486398425181e-07]

[0, 1, 2, 1, 0]
[0.5, 0.5, 1.0, 0.5, 0.5]
[12]:
# Derivatives of twiss parameters (chromaticity)

# pip install git+https://github.com/i-a-morozov/twiss.git@main
# pip install git+https://github.com/i-a-morozov/ndmap.git@main

from twiss import twiss

from ndmap.pfp import parametric_fixed_point
from ndmap.evaluate import evaluate

# Define elements

QF = Quadrupole('QF', 0.5, +0.21)
QD = Quadrupole('QD', 0.5, -0.19)
SF = Sextupole('SF', 0.25)
SD = Sextupole('SD', 0.25)
DA = Drift('DR', 0.25)
DB = Drift('DR', 4.00)

# Define one-turn transformation

def fodo(state, dp, ms):
    dp, *_ = dp
    msf, msd, *_ = ms
    state = QF(state, data={**QF.data(), **{'dp': dp}})
    state = DA(state, data={**DA.data(), **{'dp': dp}})
    state = SF(state, data={**SF.data(), **{'dp': dp, 'ms': msf}})
    state = DB(state, data={**DB.data(), **{'dp': dp}})
    state = SD(state, data={**SD.data(), **{'dp': dp, 'ms': msd}})
    state = DA(state, data={**DA.data(), **{'dp': dp}})
    state = QD(state, data={**QD.data(), **{'dp': dp}})
    state = QD(state, data={**QD.data(), **{'dp': dp}})
    state = DA(state, data={**DA.data(), **{'dp': dp}})
    state = SD(state, data={**SD.data(), **{'dp': dp, 'ms': msd}})
    state = DB(state, data={**DB.data(), **{'dp': dp}})
    state = SF(state, data={**SF.data(), **{'dp': dp, 'ms': msf}})
    state = DA(state, data={**DA.data(), **{'dp': dp}})
    state = QF(state, data={**QF.data(), **{'dp': dp}})
    return state

# Set deviation parameters

msf = torch.tensor(0.0, dtype=torch.float64)
msd = torch.tensor(0.0, dtype=torch.float64)
ms = torch.stack([msf, msd])
dp = torch.tensor([0.0], dtype=torch.float64)

# Set fixed point

fp = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)


# Compute parametrix fixed point (first order in momentum deviation)
# Note, all parameters must be vectors

pfp, *_ = parametric_fixed_point((1, ), fp, [dp], fodo, ms)

# Define transformation around fixed point

def pfp_fodo(state, dp, ms):
    return fodo(state + evaluate(pfp, [dp]), dp, ms) - evaluate(pfp, [dp])

# Tune

def tune(dp, ms):
    matrix = torch.func.jacrev(pfp_fodo)(fp, dp, ms)
    tune, *_ = twiss(matrix)
    return tune

# Chromaticity

def chromaticity(ms):
    return torch.func.jacrev(tune)(dp, ms)

# Compute tunes

tunes = tune(dp, ms)
print(tunes)

# Compute chromaticity

chromaticities = chromaticity(ms)
print(chromaticities.squeeze())

# Compute derivative of chromaticities
# The result is zero, since there is no dispersion to feed sextupoles down

print(torch.func.jacrev(chromaticity)(ms).squeeze())
tensor([0.2107, 0.1703], dtype=torch.float64)
tensor([-0.2279, -0.2107], dtype=torch.float64)
tensor([[0., 0.],
        [0., 0.]], dtype=torch.float64)

Example-08: Octupole (element)

[1]:
# Comparison of octupole element with MADX-PTC and other features
[2]:
from pathlib import Path
from os import system

import torch
from model.library.octupole import Octupole
[3]:
# Tracking (paraxial)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = False
align = False

mo = + 50.0
dp = 0.005
length = 0.2
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag: octupole, l={length},k3={mo};
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

O = Octupole('O', length=length, mo=mo, dp=dp, exact=exact, order=5, ns=10)
res = O(state, alignment=align, data={**O.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.009004942132987458, -0.005000291994759593, -0.004801204788639501, 0.0009979801079392878]
[0.009004942132987486, -0.005000291994759635, -0.004801204788639504, 0.0009979801079392843]
[-2.7755575615628914e-17, 4.2500725161431774e-17, 3.469446951953614e-18, 3.469446951953614e-18]
[4]:
# Tracking (exact)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = True
align = False

mo = + 50.0
dp = 0.005
length = 0.2
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag: octupole, l={length},k3={mo};
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

O = Octupole('O', length=length, mo=mo, dp=dp, exact=exact, order=5, ns=10)
res = O(state, alignment=align, data={**O.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.00900492932570409, -0.005000291993307186, -0.004801202229721349, 0.0009979801112567327]
[0.009004929325704207, -0.0050002919933071705, -0.004801202229721335, 0.00099798011125672]
[-1.1622647289044608e-16, -1.5612511283791264e-17, -1.3877787807814457e-17, 1.2576745200831851e-17]
[5]:
# Tracking (exact, alignment)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = True
align = True

mo = + 50.0
dp = 0.005
length = 0.2
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag: octupole, l={length},k3={mo};
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

O = Octupole('O', length=length, mo=mo, dp=dp, exact=exact, order=5, ns=10)
res = O(state, alignment=align, data={**O.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.009005660276793346, -0.004983923494470567, -0.004794436076349048, 0.0011299481578555504]
[0.009005660276792638, -0.004983923494470562, -0.004794436076348805, 0.0011299481578555233]
[7.077671781985373e-16, -5.204170427930421e-18, -2.42861286636753e-16, 2.710505431213761e-17]
[6]:
# Deviation/error variables

mo = 50.0
dp = 0.005
length = 0.2
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

O = Octupole('O', length, mo, dp)

# Each element has two variant of a call method
# In the first case only state is passed, it is transformed using parameters specified on initializaton

print(O(state))
print()

# Deviation errors can be also passed to call method
# These variables are added to corresponding parameters specified on initializaton
# For example, element lenght can changed

print(O(state, data={**O.data(), **{'dl': -O.length}}))
print()

# In the above O.data() creates default deviation dictionary (with zero values for each deviaton)
# {**O.data(), **{'dl': -O.length}} replaces the 'dl' key value

# Additionaly, alignment errors are passed as deivation variables
# They are used if alignment flag is raised

print(O(state, data={**O.data(), **error}, alignment=True))
print()

# The following elements can be made equivalent using deviation variables

OA = Octupole('OA', length, mo, dp)
OB = Octupole('OB', length - 0.1, mo, dp)

print(OA(state) - OB(state, data={**OB.data(), **{'dl': torch.tensor(+0.1, dtype=OB.dtype)}}))

# Note, while in some cases float values can be passed as values to deviation variables
# The correct behaviour in guaranteed only for tensors
tensor([ 0.0090, -0.0050, -0.0048,  0.0010], dtype=torch.float64)

tensor([ 0.0100, -0.0050, -0.0050,  0.0010], dtype=torch.float64)

tensor([ 0.0090, -0.0050, -0.0048,  0.0011], dtype=torch.float64)

tensor([0., 0., 0., 0.], dtype=torch.float64)
[7]:
# Insertion element

# In this mode elements are treated as thin insertions (at the center)
# Using parameters specified on initialization, transport two matrices are computed
# These matrices are used to insert the element
# Input state is transformed from the element center to its entrance
# Next, transformation from the entrance frame to the exit frame is performed
# This transformation can contain errors
# The final step is to transform state from the exit frame back to the element center
# Without errors, this results in identity transformation for linear elements

mo = 50.0
dp = 0.005
length = 0.2
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

O = Octupole('O', length, mo, dp, exact=False, insertion=True)

# Since octupole is a nonlinear element, insertion is an identity transformation only for zero strenght

print(O(state) - state)
print(O(state, data={**O.data(), **{'mo': -mo}}) - state)

# Represents effect of an error (any nonzero value of strengh or a change in other parameter)

print(O(state, data={**O.data(), **{'dl': 0.1}}) - state)

# Exact tracking corresponds to inclusion of kinematic term as errors

O = Octupole('O', length, mo, dp, exact=True, insertion=True, ns=100, order=1)

print(O(state) - state)
tensor([ 0.0000e+00, -4.1667e-07,  0.0000e+00, -2.2917e-06],
       dtype=torch.float64)
tensor([0., 0., 0., 0.], dtype=torch.float64)
tensor([-4.9754e-04, -5.2588e-07,  9.9342e-05, -3.2270e-06],
       dtype=torch.float64)
tensor([-1.7351e-08, -4.1976e-07, -6.9314e-09, -2.2953e-06],
       dtype=torch.float64)
[8]:
# Mapping over a set of initial conditions

# Call method can be used to map over a set of initial conditions
# Note, device can be set to cpu or gpu via base element classvariables

mo = 50.0
dp = 0.005
length = 0.2

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

O = Octupole('O', length, mo, dp, exact=True)

state = 1.0E-3*torch.randn((512, 4), dtype=O.dtype, device=O.device)

print(torch.vmap(O)(state).shape)

# To map over deviations parameters a wrapper function (or a lambda expression) can be used

def wrapper(state, dp):
    return O(state, data={**O.data(), **{'dp': dp}})

dp = 1.0E-3*torch.randn(512, dtype=O.dtype, device=O.device)

print(torch.vmap(wrapper)(state, dp).shape)
torch.Size([512, 4])
torch.Size([512, 4])
[9]:
# Differentiability

# Both call methods are differentiable
# Derivative with respect to state can be computed directly
# For deviation variables, wrapping is required

mo = 50.0
dp = 0.005
length = 0.2
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

O = Octupole('O', length, mo, dp, exact=False)

# Compute derivative with respect to state

print(torch.func.jacrev(O)(state))
print()

# Compute derivative with respect to a deviation variable

mo = torch.tensor(0.0, dtype=torch.float64)

def wrapper(state, mo):
    return O(state, data={**O.data(), **{'mo': mo}})

print(torch.func.jacrev(wrapper, 1)(state, mo))
print()
tensor([[ 9.9997e-01,  1.9900e-01, -4.6335e-05, -4.6105e-06],
        [-3.3141e-04,  9.9997e-01, -4.6567e-04, -4.6335e-05],
        [-4.6335e-05, -4.6105e-06,  1.0000e+00,  1.9901e-01],
        [-4.6567e-04, -4.6335e-05,  3.3141e-04,  1.0000e+00]],
       dtype=torch.float64)

tensor([-5.7528e-10, -5.7815e-09, -4.0127e-09, -4.0327e-08],
       dtype=torch.float64)

[10]:
# Output at each step

# It is possible to collect output of state or tangent matrix at each integration step
# Number of integratin steps is controlled by ns parameter on initialization
# Alternatively, desired integration step length can be passed
# Number of integration steps is computed as ceil(length/ds)

mo = 50.0
dp = 0.005
length = 0.2
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

O = Octupole('O', length, mo, dp, exact=False, ns=10, output=True, matrix=True)

# Final state is still returned

print(O(state))

# Data is added to special attributes (state and tangent matrix)

print(O.container_output.shape)
print(O.container_matrix.shape)

# Number of integration steps can be changed

O.ns = 100

O(state)
print(O.container_output.shape)
print(O.container_matrix.shape)
tensor([ 0.0090, -0.0050, -0.0048,  0.0010], dtype=torch.float64)
torch.Size([10, 4])
torch.Size([10, 4, 4])
torch.Size([100, 4])
torch.Size([100, 4, 4])
[11]:
# Integration order is set on initialization (default value is zero)
# This order is related to difference order as 2n + 2
# Thus, zero corresponds to second order difference method

mo = 50.0
dp = 0.005
length = 0.2
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

O = Octupole('O', length, mo, dp, order=0, exact=True)

# For octupole integration is always performed
# In exact case, kinematic term error is added

O.ns = 10
ref = O(state)

O.ns = 100
res = O(state)

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())
print()

# Integrator parameters are stored in data attribute (if integration is actually performed)

maps, weights = O._data
print(maps)
print(weights)
[0.00900492936806216, -0.005000291964120008, -0.004801202138358545, 0.0009979801465156844]
[0.00900492932612774, -0.005000291993015304, -0.004801202228807697, 0.0009979801116093128]
[4.193442121325219e-11, 2.889529587823958e-11, 9.044915164069245e-11, 3.4906371629978006e-11]

[0, 1, 2, 1, 0]
[0.5, 0.5, 1.0, 0.5, 0.5]

Example-09: Multipole (element)

[1]:
# Comparison of sextupole element with MADX-PTC and other features
[2]:
from pathlib import Path
from os import system

import torch
from model.library.multipole import Multipole
[3]:
# Tracking (paraxial)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = False
align = False

kn = - 2.0
ks = + 1.5
ms = 25.0
mo = 110.0
dp = 0.005
length = 0.5
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:quadrupole,l={length},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}};
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

M = Multipole('M', length=length, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=5, ns=10)
res = M(state, alignment=align, data={**M.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.00892946929535585, 0.0009165531774213284, -0.0017930189940976037, 0.01128986832436645]
[0.008929469295355376, 0.0009165531774212371, -0.001793018994097562, 0.011289868324366148]
[4.735795089416683e-16, 9.128982292327947e-17, -4.163336342344337e-17, 3.0184188481996443e-16]
[4]:
# Tracking (exact)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = True
align = False

kn = - 2.0
ks = + 1.5
ms = 25.0
mo = 110.0
dp = 0.005
length = 0.5
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:quadrupole,l={length},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}};
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

M = Multipole('M', length=length, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=5, ns=10)
res = M(state, alignment=align, data={**M.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.00892945162589142, 0.0009165626425029898, -0.0017929090323601758, 0.011289834389170184]
[0.008929451625891005, 0.0009165626425029993, -0.001792909032360122, 0.01128983438916985]
[4.145989107584569e-16, -9.432558900623889e-18, -5.377642775528102e-17, 3.3480163086352377e-16]
[5]:
# Tracking (exact, alignment)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = True
align = True

kn = - 2.0
ks = + 1.5
ms = 25.0
mo = 110.0
dp = 0.005
length = 0.5
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:quadrupole,l={length},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}};
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

M = Multipole('M', length=length, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=5, ns=10)
res = M(state, alignment=align, data={**M.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.001664975746516068, -0.04016261215593401, -0.015906155451718536, -0.05363979979736713]
[0.0016649757465192622, -0.04016261215593149, -0.015906155451718116, -0.053639799797364655]
[-3.194059600142296e-15, -2.518818487118324e-15, -4.198030811863873e-16, -2.4771851236948805e-15]
[6]:
# Deviation/error variables

kn = - 2.0
ks = + 1.5
ms = 25.0
mo = 110.0
mo = 50.0
dp = 0.005
length = 0.2
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

M = Multipole('M', length, kn, ks, ms, mo, dp)

# Each element has two variant of a call method
# In the first case only state is passed, it is transformed using parameters specified on initializaton

print(M(state))
print()

# Deviation errors can be also passed to call method
# These variables are added to corresponding parameters specified on initializaton
# For example, element lenght can changed

print(M(state, data={**M.data(), **{'dl': -M.length}}))
print()

# In the above M.data() creates default deviation dictionary (with zero values for each deviaton)
# {**M.data(), **{'dl': -M.length}} replaces the 'dl' key value

# Additionaly, alignment errors are passed as deivation variables
# They are used if alignment flag is raised

print(M(state, data={**M.data(), **error}, alignment=True))
print()

# The following elements can be made equivalent using deviation variables

MA = Multipole('MA', length, kn, ks, ms, mo, dp)
MB = Multipole('MB', length - 0.1, kn, ks, ms, mo, dp)

print(MA(state) - MB(state, data={**MB.data(), **{'dl': torch.tensor(+0.1, dtype=MB.dtype)}}))

# Note, while in some cases float values can be passed as values to deviation variables
# The correct behaviour in guaranteed only for tensors
tensor([ 0.0092, -0.0028, -0.0043,  0.0055], dtype=torch.float64)

tensor([ 0.0100, -0.0050, -0.0050,  0.0010], dtype=torch.float64)

tensor([ 0.0085, -0.0159, -0.0060, -0.0224], dtype=torch.float64)

tensor([0., 0., 0., 0.], dtype=torch.float64)
[7]:
# Insertion element

# In this mode elements are treated as thin insertions (at the center)
# Using parameters specified on initialization, transport two matrices are computed
# These matrices are used to insert the element
# Input state is transformed from the element center to its entrance
# Next, transformation from the entrance frame to the exit frame is performed
# This transformation can contain errors
# The final step is to transform state from the exit frame back to the element center
# Without errors, this results in identity transformation for linear elements

kn = - 2.0
ks = + 1.5
ms = 25.0
mo = 110.0
mo = 50.0
dp = 0.005
length = 0.2
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

M = Multipole('M', length, kn, ks, ms, mo, dp, exact=False, insertion=True)

# Since multipole is a nonlinear element (non-zero sextupole or octupole)
# Insertion is an identity transformation only for zero strenght

print(M(state) - state)
print(M(state, data={**M.data(), **{'ms': -ms, 'mo': -mo}}) - state)

# Represents effect of an error (any nonzero value of strengh or a change in other parameter)

print(M(state, data={**M.data(), **{'dl': 0.1}}) - state)

# Exact tracking corresponds to inclusion of kinematic term as errors

M = Multipole('M', length, kn, ks, ms, mo, dp, exact=True, insertion=True, ns=20, order=1)

print(M(state) - state)
tensor([ 6.9389e-18, -1.8792e-04,  8.6736e-19, -2.5229e-04],
       dtype=torch.float64)
tensor([ 6.9389e-18, -4.3368e-18,  8.6736e-19, -8.6736e-19],
       dtype=torch.float64)
tensor([-0.0004,  0.0009,  0.0002,  0.0021], dtype=torch.float64)
tensor([-7.9785e-07, -1.9093e-04, -5.6642e-07, -2.5071e-04],
       dtype=torch.float64)
[8]:
# Mapping over a set of initial conditions

# Call method can be used to map over a set of initial conditions
# Note, device can be set to cpu or gpu via base element classvariables

kn = - 2.0
ks = + 1.5
ms = 25.0
mo = 110.0
mo = 50.0
dp = 0.005
length = 0.2

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

M = Multipole('M', length, kn, ks, ms, mo, dp, exact=True)

state = 1.0E-3*torch.randn((512, 4), dtype=M.dtype, device=M.device)

print(torch.vmap(M)(state).shape)

# To map over deviations parameters a wrapper function (or a lambda expression) can be used

def wrapper(state, dp):
    return M(state, data={**M.data(), **{'dp': dp}})

dp = 1.0E-3*torch.randn(512, dtype=M.dtype, device=M.device)

print(torch.vmap(wrapper)(state, dp).shape)
torch.Size([512, 4])
torch.Size([512, 4])
[9]:
# Differentiability

# Both call methods are differentiable
# Derivative with respect to state can be computed directly
# For deviation variables, wrapping is required

kn = - 2.0
ks = + 1.5
ms = 25.0
mo = 110.0
mo = 50.0
dp = 0.005
length = 0.2
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

M = Multipole('M', length, kn, ks, ms, mo, dp, exact=False)

# Compute derivative with respect to state

print(torch.func.jacrev(M)(state))
print()

# Compute derivative with respect to a deviation variable

kn = torch.tensor(0.0, dtype=torch.float64)

def wrapper(state, kn):
    return M(state, data={**M.data(), **{'kn': kn}})

print(torch.func.jacrev(wrapper, 1)(state, kn))
print()
tensor([[ 1.0353,  0.2012,  0.0274,  0.0017],
        [ 0.3588,  1.0353,  0.2757,  0.0274],
        [ 0.0274,  0.0017,  0.9653,  0.1969],
        [ 0.2757,  0.0274, -0.3449,  0.9653]], dtype=torch.float64)

tensor([-1.9468e-04, -1.9487e-03, -9.6920e-05, -9.5528e-04],
       dtype=torch.float64)

[10]:
# Output at each step

# It is possible to collect output of state or tangent matrix at each integration step
# Number of integratin steps is controlled by ns parameter on initialization
# Alternatively, desired integration step length can be passed
# Number of integration steps is computed as ceil(length/ds)

kn = - 2.0
ks = + 1.5
ms = 25.0
mo = 110.0
mo = 50.0
dp = 0.005
length = 0.2
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

M = Multipole('M', length, kn, ks, ms, mo, dp, exact=False, ns=10, output=True, matrix=True)

# Final state is still returned

print(M(state))

# Data is added to special attributes (state and tangent matrix)

print(M.container_output.shape)
print(M.container_matrix.shape)

# Number of integration steps can be changed

M.ns = 100

M(state)
print(M.container_output.shape)
print(M.container_matrix.shape)
tensor([ 0.0092, -0.0028, -0.0043,  0.0055], dtype=torch.float64)
torch.Size([10, 4])
torch.Size([10, 4, 4])
torch.Size([100, 4])
torch.Size([100, 4, 4])
[11]:
# Integration order is set on initialization (default value is zero)
# This order is related to difference order as 2n + 2
# Thus, zero corresponds to second order difference method

kn = - 2.0
ks = + 1.5
ms = 25.0
mo = 110.0
mo = 50.0
dp = 0.005
length = 0.2
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

M = Multipole('M', length, kn, ks, ms, mo, dp, order=0, exact=True)

# For multipole with non-zero sextupole and/or octupole integration is always performed
# In exact case, kinematic term error is added

M.ns = 10
ref = M(state)

M.ns = 100
res = M(state)

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())
print()

# Integrator parameters are stored in data attribute (if integration is actually performed)

maps, weights = M._data
print(maps)
print(weights)
[0.009228705737231328, -0.002766282719416843, -0.004341639948599635, 0.005542200849934217]
[0.009228699312521686, -0.0027663108074714475, -0.004341646792491436, 0.00554221421387834]
[6.424709642766091e-09, 2.808805460441377e-08, 6.843891801368296e-09, -1.3363944122331273e-08]

[0, 1, 2, 1, 0]
[0.5, 0.5, 1.0, 0.5, 0.5]

Example-10: Dipole (element)

[1]:
# Comparison of dipole element with MADX-PTC and other features

# Note, cylindrical multipole is included upto octupole order
# Potential is not truncated in paraxial case, only sqrt is expanded
# In exact case effects of multipoles are not accounted in wedges (cases with e1 or e2 not equal to zero)
[2]:
from pathlib import Path
from os import system

import torch
from model.library.drift import Drift
from model.library.quadrupole import Quadrupole
from model.library.sextupole import Sextupole
from model.library.dipole import Dipole
[3]:
# Tracking (paraxial)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = False
align = False

length = 2.0
angle = 0.05
e1 = 0.0
e2 = 0.0
kn = 0.0
ks = 0.0
ms = 0.0
mo = 0.0
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[5.160257777686003e-05, -0.004956273150572646, -0.0030019980019983517, 0.001]
[5.1602577776445975e-05, -0.004956273150572915, -0.0030019980019981656, 0.0010000000000000434]
[4.1405680967221414e-16, 2.688821387764051e-16, -1.8604909279851256e-16, -4.336808689942018e-17]
[4]:
# Tracking (paraxial, e1 & e2)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = False
align = False

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
kn = 0.0
ks = 0.0
ms = 0.0
mo = 0.0
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[6.408749411377232e-05, -0.00494998958983017, -0.002995752944647053, 0.0010049983869644134]
[6.408749411335537e-05, -0.004949989589830473, -0.0029957529446465857, 0.0010049983869644566]
[4.169434979564568e-16, 3.0270924655795284e-16, -4.670742959067553e-16, -4.3151246464923076e-17]
[5]:
# Tracking (paraxial, e1 & e2, alignment)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = False
align = True

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
kn = 0.0
ks = 0.0
ms = 0.0
mo = 0.0
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.0027250502286323983, -0.004697460778410744, -0.007781317874336121, -0.004015687563474919]
[0.0027250502286311207, -0.004697460778410097, -0.007781317874338802, -0.004015687563476177]
[1.2776238400569184e-15, -6.47051856539349e-16, 2.6810151321221554e-15, 1.2576745200831851e-15]
[6]:
# Tracking (paraxial, kn & ks)

# Note, in paraxial case, MADX seems to truncate cylindrical potential

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = False
align = False

length = 2.0
angle = 0.05
e1 = 0.0
e2 = 0.0
kn = 0.1
ks = 0.1
ms = 0.0
mo = 0.0
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[-0.002043638956596164, -0.006566087436123452, -0.00259447159458907, 0.0010855376916213359]
[-0.0020440852924812135, -0.006566346702412346, -0.002594440402484201, 0.001085525276668085]
[4.46335885049675e-07, 2.592662888935629e-07, -3.119210486906762e-08, 1.2414953250768773e-08]
[7]:
# Tracking (paraxial, kn & ks, alignment)

# Note, in paraxial case, MADX seems to truncate cylindrical potential

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = False
align = True

length = 2.0
angle = 0.05
e1 = 0.0
e2 = 0.0
kn = 0.1
ks = 0.1
ms = 0.0
mo = 0.0
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.013167193439103305, 0.005511119368028012, -0.010781508233295324, -0.007329921037793256]
[0.013155910631910853, 0.0055012372192967465, -0.010779780613239438, -0.00732819304969001]
[1.128280719245138e-05, 9.882148731265271e-06, -1.727620055886822e-06, -1.7279881032459046e-06]
[8]:
# Tracking (paraxial, kn & ks, ms & mo, alignment)

# Note, in paraxial case, MADX seems to truncate cylindrical potential

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = False
align = True

length = 2.0
angle = 0.05
e1 = 0.0
e2 = 0.0
kn = 0.1
ks = 0.1
ms = 0.1
mo = 0.1
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.013072849472264409, 0.005416718324594963, -0.010908836090706359, -0.007452351251381879]
[0.013061618858281962, 0.0054068762860350235, -0.010906993838412838, -0.007450515116390934]
[1.1230613982447443e-05, 9.842038559939453e-06, -1.842252293521307e-06, -1.8361349909453567e-06]
[9]:
# Tracking (paraxial, e1 & e2, kn & ks, ms & mo, alignment)

# Note, in paraxial case, MADX seems to truncate cylindrical potential
# Note, model dipole doesn't account for multipoles in wedges

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = False
align = True

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
kn = 0.1
ks = 0.1
ms = 0.1
mo = 0.1
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.013030972087034147, 0.005379072261796265, -0.010940636831296861, -0.007479587493733525]
[0.013019736470996354, 0.005369216085022485, -0.0109387887264908, -0.007477745021700929]
[1.1235616037793758e-05, 9.85617677377957e-06, -1.8481048060618732e-06, -1.8424720325954658e-06]
[10]:
# Tracking (exact, alignment)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = True
align = True

length = 2.0
angle = 0.05
e1 = 0.0
e2 = 0.0
kn = 0.0
ks = 0.0
ms = 0.0
mo = 0.0
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.002770798305588665, -0.004651371270674763, -0.007745680488823617, -0.0039921009984161026]
[0.0027707983054956944, -0.004651371270673946, -0.007745680488833436, -0.003992100998417344]
[9.2970769971501e-14, -8.161873954470877e-16, 9.819402235766717e-15, 1.2411946470614055e-15]
[11]:
# Tracking (exact, kn & ks, ms & mo, alignment)

# Note, in model dipole cylindrical potential is truncated at octupole

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = True
align = True

length = 2.0
angle = 0.05
e1 = 0.0
e2 = 0.0
kn = 0.1
ks = 0.1
ms = 0.1
mo = 0.1
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.013055723507830729, 0.005405610520344896, -0.01089583005089035, -0.007447552679263773]
[0.013055723646617522, 0.005405610850033292, -0.010895828250285611, -0.007447551055280752]
[-1.387867928220876e-10, -3.2968839637492753e-10, -1.8006047382279622e-09, -1.6239830209763273e-09]
[12]:
# Tracking (exact, e1 & e2, kn & ks, ms & mo, alignment)

# Note, in model dipole cylindrical potential is truncated at octupole
# Note, model dipole doesn't account for multipoles in wedges

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = True
align = True

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
kn = 0.1
ks = 0.1
ms = 0.1
mo = 0.1
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag:sbend,l={length},angle={angle},e1={e1},e2={e2},knl={{0.0,{kn*length},{ms*length},{mo*length}}},ksl={{0.0,{ks*length}}},kill_ent_fringe=false,kill_exi_fringe=false;
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=5,sector_nmul=5 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

D = Dipole('D', length=length, angle=angle, e1=e1, e2=e2, kn=kn, ks=ks, ms=ms, mo=mo, dp=dp, exact=exact, order=2, ns=25)
res = D(state, alignment=align, data={**D.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.013018589129527029, 0.005372361621197678, -0.010919187705983761, -0.007468383469439608]
[0.013013881710648583, 0.0053679558070361885, -0.010927596957624465, -0.007474796553899775]
[4.707418878445446e-06, 4.4058141614898225e-06, 8.409251640703955e-06, 6.413084460166543e-06]
[13]:
# Deviation/error variables

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
kn = 0.1
ks = 0.1
ms = 0.1
mo = 0.1
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp)

# Each element has two variant of a call method
# In the first case only state is passed, it is transformed using parameters specified on initializaton

print(D(state))
print()

# Deviation errors can be also passed to call method
# These variables are added to corresponding parameters specified on initializaton
# For example, element lenght can changed

print(D(state, data={**D.data(), **{'dl': -0.5*D.length}}))
print()

# In the above D.data() creates default deviation dictionary (with zero values for each deviaton)
# {**D.data(), **{'dl': -D.length}} replaces the 'dl' key value

# Additionaly, alignment errors are passed as deivation variables
# They are used if alignment flag is raised

print(D(state, data={**D.data(), **error}, alignment=True))
print()

# The following elements can be made equivalent using deviation variables

DA = Dipole('DA', length, angle, e1, e2, kn, ks, ms, mo, dp)
DB = Dipole('DB', length - 0.1, angle, e1, e2, kn, ks, ms, mo, dp)

print(DA(state) - DB(state, data={**DB.data(), **{'dl': torch.tensor(+0.1, dtype=DB.dtype)}}))

# Note, while in some cases float values can be passed as values to deviation variables
# The correct behaviour in guaranteed only for tensors
tensor([-0.0020, -0.0066, -0.0026,  0.0011], dtype=torch.float64)

tensor([ 0.0044, -0.0061, -0.0038,  0.0013], dtype=torch.float64)

tensor([ 0.0130,  0.0054, -0.0109, -0.0075], dtype=torch.float64)

tensor([0., 0., 0., 0.], dtype=torch.float64)
[14]:
# Insertion element

import torch
from model.library.dipole import Dipole

# In this mode elements are treated as thin insertions (at the center)
# Using parameters specified on initialization, transport two matrices are computed
# These matrices are used to insert the element
# Input state is transformed from the element center to its entrance
# Next, transformation from the entrance frame to the exit frame is performed
# This transformation can contain errors
# The final step is to transform state from the exit frame back to the element center
# Without errors, this results in identity transformation for linear elements

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
kn = 0.0
ks = 0.0
ms = 0.0
mo = 0.0
dp = 0.0
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp, exact=False, insertion=True)

# For dipole insertion is an identity transformation in the following is true
# dp = 0 (chomatic orbit change)
# kn & ks and ms & mo are all equal to zero (nonlinearity of cylindrical potential)

print(D(state) - state)

# Represents effect of an error (any nonzero value of strengh or a change in other parameter)

print(D(state, data={**D.data(), **{'dp': 0.001}}) - state)

# Exact tracking corresponds to inclusion of kinematic term as errors

D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp, exact=True, insertion=True)

print(D(state) - state)
tensor([ 1.0408e-17, -4.3368e-18, -4.3368e-18,  8.6736e-19],
       dtype=torch.float64)
tensor([ 9.9380e-06,  4.9995e-05, -1.9980e-06,  6.5052e-19],
       dtype=torch.float64)
tensor([-2.3859e-06, -5.8600e-06, -7.2738e-07, -2.4947e-07],
       dtype=torch.float64)
[15]:
# Mapping over a set of initial conditions

# Call method can be used to map over a set of initial conditions
# Note, device can be set to cpu or gpu via base element classvariables

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
kn = 0.5
ks = 0.5
ms = 1.0
mo = 1.0
dp = 0.001

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp)

state = 1.0E-3*torch.randn((512, 4), dtype=D.dtype, device=D.device)

print(torch.vmap(D)(state).shape)

# To map over deviations parameters a wrapper function (or a lambda expression) can be used

def wrapper(state, dp):
    return D(state, data={**D.data(), **{'dp': dp}})

dp = 1.0E-3*torch.randn(512, dtype=D.dtype, device=D.device)

print(torch.vmap(wrapper)(state, dp).shape)
torch.Size([512, 4])
torch.Size([512, 4])
[16]:
# Differentiability

# Both call methods are differentiable
# Derivative with respect to state can be computed directly
# For deviation variables, wrapping is required

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
kn = 0.5
ks = 0.5
ms = 1.0
mo = 1.0
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp)

# Compute derivative with respect to state

print(torch.func.jacrev(D)(state))
print()

# Compute derivative with respect to a deviation variable

e1 = torch.tensor(0.0, dtype=torch.float64)

def wrapper(state, e1):
    return D(state, data={**D.data(), **{'e1': e1}})

print(torch.func.jacrev(wrapper, 1)(state, e1))
print()
tensor([[ 0.3083,  1.4545,  1.0125,  0.6652],
        [-0.3993,  0.3083,  1.0602,  1.0168],
        [ 1.0168,  0.6652,  2.3568,  2.8065],
        [ 1.0602,  1.0125,  1.7383,  2.3568]], dtype=torch.float64)

tensor([0.0004, 0.0002, 0.0005, 0.0005], dtype=torch.float64)

[17]:
# Output at each step

# It is possible to collect output of state or tangent matrix at each integration step
# Number of integratin steps is controlled by ns parameter on initialization
# Alternatively, desired integration step length can be passed
# Number of integration steps is computed as ceil(length/ds)

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
kn = 0.5
ks = 0.5
ms = 1.0
mo = 1.0
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp, exact=False, ns=10, output=True, matrix=True)

# Final state is still returned

print(D(state))

# Data is added to special attributes (state and tangent matrix)

print(D.container_output.shape)
print(D.container_matrix.shape)

# Number of integration steps can be changed

D.ns = 100

D(state)
print(D.container_output.shape)
print(D.container_matrix.shape)
tensor([-0.0086, -0.0098, -0.0022, -0.0008], dtype=torch.float64)
torch.Size([10, 4])
torch.Size([10, 4, 4])
torch.Size([100, 4])
torch.Size([100, 4, 4])
[18]:
# Integration order is set on initialization (default value is zero)
# This order is related to difference order as 2n + 2
# Thus, zero corresponds to second order difference method

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
kn = 0.5
ks = 0.5
ms = 1.0
mo = 1.0
dp = 0.001
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

D = Dipole('D', length, angle, e1, e2, kn, ks, ms, mo, dp, exact=True, ns=10)

# For dipole integration is always performed
# In exact case, kinematic term error is added

D.ns = 10
ref = D(state)

D.ns = 100
res = D(state)

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())
print()

# Integrator parameters are stored in data attribute (if integration is actually performed)

maps, weights = D._data
print(maps)
print(weights)
[-0.00858751580977462, -0.009819760091907982, -0.0021662845096666727, -0.0008078063877268448]
[-0.008587660168877631, -0.00982000432362944, -0.0021667806526843485, -0.0008081645802833406]
[1.443591030117486e-07, 2.442317214579326e-07, 4.961430176758305e-07, 3.5819255649581234e-07]

[0, 1, 2, 3, 4, 3, 2, 1, 0]
[0.5, 0.5, 0.5, 0.5, 1.0, 0.5, 0.5, 0.5, 0.5]
[19]:
# Derivatives of twiss parameters (chromaticity)

# pip install git+https://github.com/i-a-morozov/twiss.git@main
# pip install git+https://github.com/i-a-morozov/ndmap.git@main

from twiss import twiss

from ndmap.pfp import parametric_fixed_point
from ndmap.evaluate import evaluate

QF = Quadrupole('QF', 0.5, +0.21)
QD = Quadrupole('QD', 0.5, -0.19)
SF = Sextupole('SF', 0.25)
SD = Sextupole('SD', 0.25)
DR = Drift('DR', 0.25)
BM = Dipole('DR', 3.50, 0.15)

def fodo(state, dp, ms):
    dp, *_ = dp
    msf, msd, *_ = ms
    state = QF(state, data={**QF.data(), **{'dp': dp}})
    state = DR(state, data={**DR.data(), **{'dp': dp}})
    state = SF(state, data={**SF.data(), **{'dp': dp, 'ms': msf}})
    state = DR(state, data={**DR.data(), **{'dp': dp}})
    state = BM(state, data={**BM.data(), **{'dp': dp}})
    state = DR(state, data={**DR.data(), **{'dp': dp}})
    state = SD(state, data={**SD.data(), **{'dp': dp, 'ms': msd}})
    state = DR(state, data={**DR.data(), **{'dp': dp}})
    state = QD(state, data={**QD.data(), **{'dp': dp}})
    state = QD(state, data={**QD.data(), **{'dp': dp}})
    state = DR(state, data={**DR.data(), **{'dp': dp}})
    state = SD(state, data={**SD.data(), **{'dp': dp, 'ms': msd}})
    state = DR(state, data={**DR.data(), **{'dp': dp}})
    state = BM(state, data={**BM.data(), **{'dp': dp}})
    state = DR(state, data={**DR.data(), **{'dp': dp}})
    state = SF(state, data={**SF.data(), **{'dp': dp, 'ms': msf}})
    state = DR(state, data={**DR.data(), **{'dp': dp}})
    state = QF(state, data={**QF.data(), **{'dp': dp}})
    return state

# Set deviation parameters

msf = torch.tensor(0.0, dtype=torch.float64)
msd = torch.tensor(0.0, dtype=torch.float64)
ms = torch.stack([msf, msd])
dp = torch.tensor([0.0], dtype=torch.float64)

# Set fixed point

fp = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)


# Compute parametrix fixed point (first order in momentum deviation)
# Note, all parameters must be vectors

pfp, *_ = parametric_fixed_point((1, ), fp, [dp], fodo, ms)

# Define transformation around fixed point

def pfp_fodo(state, dp, ms):
    return fodo(state + evaluate(pfp, [dp]), dp, ms) - evaluate(pfp, [dp])

# Tune

def tune(dp, ms):
    matrix = torch.func.jacrev(pfp_fodo)(fp, dp, ms)
    tune, *_ = twiss(matrix)
    return tune

# Chromaticity

def chromaticity(ms):
    return torch.func.jacrev(tune)(dp, ms)

# Compute tunes

tunes = tune(dp, ms)
print(tunes)


# Compute chromaticity

chromaticities = chromaticity(ms).squeeze()
print(chromaticities)

# Compute derivative of chromaticities

jacobian = torch.func.jacrev(chromaticity)(ms).squeeze()
print(jacobian)

# Correct chomaticity

print((chromaticity(ms - torch.linalg.pinv(jacobian) @ chromaticities)).squeeze())
tensor([0.2210, 0.1703], dtype=torch.float64)
tensor([-0.2310, -0.2107], dtype=torch.float64)
tensor([[ 1.6013,  0.3613],
        [-0.7593, -1.2452]], dtype=torch.float64)
tensor([-1.1102e-16,  1.3878e-16], dtype=torch.float64)

Example-11: Corrector (element)

[1]:
# Comparison of corrector element with MADX-PTC and other features
[2]:
from pathlib import Path
from os import system

import torch
from model.library.corrector import Corrector
[3]:
# Tracking

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = True
align = False

cx = +0.001
cy = -0.005

dp = 0.005

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
cx:hkicker,l=0.0,kick={cx};
cy:vkicker,l=0.0,kick={cy};
map:line=(cx, cy) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

C = Corrector('C', cx=cx, cy=cy, dp=dp)
res = C(state, alignment=align, data={**C.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.0, 0.001, 0.0, -0.005]
[0.0, 0.001, 0.0, -0.005]
[0.0, 0.0, 0.0, 0.0]
[4]:
# Tracking (alignment)

# Only dx and dy alignment errors seems to work as expected in MADX
# Also, wz rotation matches tilt

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = True
align = True

cx = +0.001
cy = -0.005
dp = 0.005
state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.0, dtype=torch.float64)

wx = align*torch.tensor(0.0, dtype=torch.float64)
wy = align*torch.tensor(0.0, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
cx:hkicker,l=0.0,kick={cx},tilt={wz.item()};
cy:vkicker,l=0.0,kick={cy},tilt={wz.item()};
map:line=(cx, cy) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()} ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

C = Corrector('C', cx=cx, cy=cy, dp=dp)
res = C(state, alignment=align, data={**C.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.0, 0.0014941712485121669, 0.0, -0.004875187409743301]
[-6.938893903907228e-18, 0.0014941712485121667, 0.0, -0.004875187409743301]
[6.938893903907228e-18, 2.168404344971009e-19, 0.0, 0.0]
[5]:
# Deviation/error variables

cx = +0.001
cy = -0.005
dp = 0.005
state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)

dx = torch.tensor(+0.01, dtype=torch.float64)
dy = torch.tensor(-0.01, dtype=torch.float64)
dz = torch.tensor(0.0, dtype=torch.float64)

wx = torch.tensor(0.0, dtype=torch.float64)
wy = torch.tensor(0.0, dtype=torch.float64)
wz = torch.tensor(torch.pi, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

C = Corrector('C', cx, cy, dp)

# Each element has two variant of a call method
# In the first case only state is passed, it is transformed using parameters specified on initializaton

print(C(state))
print()

# Deviation errors can be also passed to call method
# These variables are added to corresponding parameters specified on initializaton
# For example, element lenght can changed

print(C(state, data={**C.data(), **{'cx': -cx, 'cy': -cy}}))
print()

# In the above C.data() creates default deviation dictionary (with zero values for each deviaton)
# {**C.data(), **{'cx': -cx, 'cy': -cy}} replaces the 'cx' and 'cy' key values

# Additionaly, alignment errors are passed as deivation variables
# They are used if alignment flag is raised

print(C(state, data={**C.data(), **error}, alignment=True))
print()


# The following elements can be made equivalent using deviation variables

CA = Corrector('CA', cx, cy, dp)
CB = Corrector('CB', cx - 0.001, cy, dp)

print(CA(state) - CB(state, data={**CB.data(), **{'cx': + 0.001}}))
tensor([ 0.0000,  0.0010,  0.0000, -0.0050], dtype=torch.float64)

tensor([0., 0., 0., 0.], dtype=torch.float64)

tensor([ 0.0000, -0.0010,  0.0000,  0.0050], dtype=torch.float64)

tensor([0., 0., 0., 0.], dtype=torch.float64)
[6]:
# Mapping over a set of initial conditions

# Call method can be used to map over a set of initial conditions
# Note, device can be set to cpu or gpu via base element classvariables

cx = +0.001
cy = -0.005
dp = 0.005

dx = torch.tensor(+0.01, dtype=torch.float64)
dy = torch.tensor(-0.01, dtype=torch.float64)
dz = torch.tensor(0.0, dtype=torch.float64)

wx = torch.tensor(0.0, dtype=torch.float64)
wy = torch.tensor(0.0, dtype=torch.float64)
wz = torch.tensor(torch.pi, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

C = Corrector('C', cx, cy, dp)

state = 1.0E-3*torch.randn((512, 4), dtype=C.dtype, device=C.device)

print(torch.vmap(C)(state).shape)

# To map over deviations parameters a wrapper function (or a lambda expression) can be used

def wrapper(state, cx, cy):
    return C(state, data={**C.data(), **{'cx': cx, 'cy': cy}})

cx = 1.0E-3*torch.randn(512, dtype=C.dtype, device=C.device)
cy = 1.0E-3*torch.randn(512, dtype=C.dtype, device=C.device)

print(torch.vmap(wrapper)(state, cx, cy).shape)
torch.Size([512, 4])
torch.Size([512, 4])
[7]:
# Differentiability

# Both call methods are differentiable
# Derivative with respect to state can be computed directly
# For deviation variables, wrapping is required

cx = +0.001
cy = -0.005
dp = 0.005
state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

dx = torch.tensor(+0.01, dtype=torch.float64)
dy = torch.tensor(-0.01, dtype=torch.float64)
dz = torch.tensor(0.0, dtype=torch.float64)

wx = torch.tensor(0.0, dtype=torch.float64)
wy = torch.tensor(0.0, dtype=torch.float64)
wz = torch.tensor(torch.pi, dtype=torch.float64)
error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

C = Corrector('C', cx, cy, dp)

# Compute derivative with respect to state

print(torch.func.jacrev(C)(state))
print()

# Compute derivative with respect to a deviation variable

dcx = torch.tensor(0.0, dtype=torch.float64)
dcy = torch.tensor(0.0, dtype=torch.float64)
dc = torch.stack([dcx, dcy])

def wrapper(state, dc):
    dcx, dcy = dc
    return C(state, data={**C.data(), **{'cx': dcx, 'cy': dcy}})

print(torch.func.jacrev(wrapper, 1)(state, dc))
print()
tensor([[1., 0., 0., 0.],
        [0., 1., 0., 0.],
        [0., 0., 1., 0.],
        [0., 0., 0., 1.]], dtype=torch.float64)

tensor([[0., 0.],
        [1., 0.],
        [0., 0.],
        [0., 1.]], dtype=torch.float64)

Example-12: Gradient (element)

[1]:
# Comparison of gradient element with MADX-PTC and other features
[2]:
from pathlib import Path
from os import system

import torch
from model.library.gradient import Gradient
[3]:
# Tracking

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = True
align = False

kn = +5.0E-3
ks = -2.5E-3

dp = 0.005

state = torch.tensor([0.001, 0.0, -0.005, 0.0], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag: multipole,knl={{0.0,{kn}}},ksl={{0.0,{ks}}};
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

G = Gradient('G', kn=kn, ks=ks, dp=dp)
res = G(state, alignment=align, data={**G.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.001, 7.5e-06, -0.005, -2.75e-05]
[0.001, 7.5e-06, -0.005, -2.75e-05]
[0.0, 0.0, 0.0, 0.0]
[4]:
# Tracking (alignment)

ptc = Path('ptc')
obs = Path('track.obs0001.p0001')

exact = True
align = True

kn = +5.0E-3
ks = -2.5E-3

dp = 0.005

state = torch.tensor([0.001, 0.0, -0.005, 0.0], dtype=torch.float64)
qx, px, qy, py = state.tolist()

dx = align*torch.tensor(0.05, dtype=torch.float64)
dy = align*torch.tensor(-0.02, dtype=torch.float64)
dz = align*torch.tensor(0.05, dtype=torch.float64)

wx = align*torch.tensor(0.005, dtype=torch.float64)
wy = align*torch.tensor(-0.005, dtype=torch.float64)
wz = align*torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

code = f"""
mag: multipole,knl={{0.0,{kn}}},ksl={{0.0,{ks}}};
map:line=(mag) ;
beam,energy=1.0E+6,particle=electron ;
set,format="20.20f","-20s" ;
use,period=map ;
select,flag=error,pattern="mag" ;
ealign,dx={dx.item()},dy={dy.item()},ds={dz.item()},dphi={wx.item()},dtheta={wy.item()},dpsi={wz.item()} ;
ptc_create_universe,sector_nmul_max=10,sector_nmul=10 ;
ptc_create_layout,model=1,method=6,nst=1000,exact={str(exact).lower()} ;
ptc_setswitch,fringe=false,time=true,totalpath=true,exact_mis=true ;
ptc_align ;
ptc_start,x={qx},px={px},y={qy},py={py},pt={dp},t=0.0 ;
ptc_track,icase=5,deltap=0.,turns=1,file=track,maxaper={{1.,1.,1.,1.,1.,1.}} ;
ptc_track_end ;
ptc_end ;
"""

with ptc.open('w') as stream:
    stream.write(code)

system(f'madx < {str(ptc)} > /dev/null')

with obs.open('r') as stream:
    for line in stream:
        continue
    _, _, qx, px, qy, py, *_ = line.split()

ref = torch.tensor([float(x) for x in (qx, px, qy, py)], dtype=torch.float64)

G = Gradient('G', kn=kn, ks=ks, dp=dp)
res = G(state, alignment=align, data={**G.data(), **error})

print(ref.tolist())
print(res.tolist())
print((ref - res).tolist())

ptc.unlink()
obs.unlink()
[0.000991887042015914, 0.00016412084104346932, -0.005011606734840507, 0.00023479809553215934]
[0.0009918870420159223, 0.00016412084104346948, -0.005011606734840506, 0.00023479809553215904]
[-8.239936510889834e-18, -1.6263032587282567e-19, -8.673617379884035e-19, 2.981555974335137e-19]
[5]:
# Deviation/error variables

kn = +5.0E-3
ks = -2.5E-3
dp = 0.005
state = torch.tensor([0.001, 0.0, -0.005, 0.0], dtype=torch.float64)

dx = torch.tensor(+0.01, dtype=torch.float64)
dy = torch.tensor(-0.01, dtype=torch.float64)
dz = torch.tensor(0.0, dtype=torch.float64)

wx = torch.tensor(0.0, dtype=torch.float64)
wy = torch.tensor(0.0, dtype=torch.float64)
wz = torch.tensor(torch.pi, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

G = Gradient('G', kn, ks, dp)

# Each element has two variant of a call method
# In the first case only state is passed, it is transformed using parameters specified on initializaton

print(G(state))
print()

# Deviation errors can be also passed to call method
# These variables are added to corresponding parameters specified on initializaton
# For example, element lenght can changed

print(G(state, data={**G.data(), **{'kn': -kn, 'ks': -ks}}))
print()

# In the above G.data() creates default deviation dictionary (with zero values for each deviaton)
# {**G.data(), **{'kn': -kn, 'ks': -ks}} replaces the 'kn' and 'ks' key values

# Additionaly, alignment errors are passed as deivation variables
# They are used if alignment flag is raised

print(G(state, data={**G.data(), **error}, alignment=True))
print()


# The following elements can be made equivalent using deviation variables

GA = Gradient('GA', kn, ks, dp)
GB = Gradient('GB', kn - 0.001, ks, dp)

print(GA(state) - GB(state, data={**GB.data(), **{'kn': + 0.001}}))
tensor([ 1.0000e-03,  7.5000e-06, -5.0000e-03, -2.7500e-05],
       dtype=torch.float64)

tensor([ 0.0010,  0.0000, -0.0050,  0.0000], dtype=torch.float64)

tensor([ 1.0000e-03,  3.2500e-05, -5.0000e-03,  4.7500e-05],
       dtype=torch.float64)

tensor([0., 0., 0., 0.], dtype=torch.float64)
[6]:
# Mapping over a set of initial conditions

# Call method can be used to map over a set of initial conditions
# Note, device can be set to cpu or gpu via base element classvariables

kn = +5.0E-3
ks = -2.5E-3
dp = 0.005

dx = torch.tensor(+0.01, dtype=torch.float64)
dy = torch.tensor(-0.01, dtype=torch.float64)
dz = torch.tensor(0.0, dtype=torch.float64)

wx = torch.tensor(0.0, dtype=torch.float64)
wy = torch.tensor(0.0, dtype=torch.float64)
wz = torch.tensor(torch.pi, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

G = Gradient('G', kn, ks, dp)

state = 1.0E-3*torch.randn((512, 4), dtype=G.dtype, device=G.device)

print(torch.vmap(G)(state).shape)

# To map over deviations parameters a wrapper function (or a lambda expression) can be used

def wrapper(state, kn, ks):
    return G(state, data={**G.data(), **{'kn': kn, 'ks': ks}})

kn = 1.0E-3*torch.randn(512, dtype=G.dtype, device=G.device)
ks = 1.0E-3*torch.randn(512, dtype=G.dtype, device=G.device)

print(torch.vmap(wrapper)(state, kn, ks).shape)
torch.Size([512, 4])
torch.Size([512, 4])
[7]:
# Differentiability

# Both call methods are differentiable
# Derivative with respect to state can be computed directly
# For deviation variables, wrapping is required

kn = +5.0E-3
ks = -2.5E-3
dp = 0.005
state = torch.tensor([0.001, 0.0, -0.005, 0.0], dtype=torch.float64)

dx = torch.tensor(+0.01, dtype=torch.float64)
dy = torch.tensor(-0.01, dtype=torch.float64)
dz = torch.tensor(0.0, dtype=torch.float64)

wx = torch.tensor(0.0, dtype=torch.float64)
wy = torch.tensor(0.0, dtype=torch.float64)
wz = torch.tensor(torch.pi, dtype=torch.float64)
error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

G = Gradient('G', kn, ks, dp)

# Compute derivative with respect to state

print(torch.func.jacrev(G)(state))
print()

# Compute derivative with respect to a deviation variable

dkn = torch.tensor(0.0, dtype=torch.float64)
dks = torch.tensor(0.0, dtype=torch.float64)
dk = torch.stack([dkn, dks])

def wrapper(state, dk):
    dkn, dks = dk
    return G(state, data={**G.data(), **{'kn': dkn, 'ks': dks}})

print(torch.func.jacrev(wrapper, 1)(state, dk))
print()
tensor([[ 1.0000,  0.0000,  0.0000,  0.0000],
        [-0.0050,  1.0000, -0.0025,  0.0000],
        [ 0.0000,  0.0000,  1.0000,  0.0000],
        [-0.0025,  0.0000,  0.0050,  1.0000]], dtype=torch.float64)

tensor([[-0.0000,  0.0000],
        [-0.0010, -0.0050],
        [-0.0000,  0.0000],
        [-0.0050,  0.0010]], dtype=torch.float64)

Example-13: Linear (element)

[1]:
# Can be used to model elements as linear transformations
# Constant offset can be used to model first order dispersion
[2]:
import torch
from model.library.quadrupole import Quadrupole
from model.library.dipole import Dipole
from model.library.linear import Linear
[3]:
# Linear quadrupole

length = 1.0
kn = - 2.0
ks = + 1.5
dp = 0.0
Q = Quadrupole('Q', length, kn, ks, dp)

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)

matrix = torch.func.jacrev(lambda state, dp: Q(state, data={**Q.data(), **{'dp': dp}}), 0)(state, Q.dp)
vector = torch.func.jacrev(lambda state, dp: Q(state, data={**Q.data(), **{'dp': dp}}), 1)(state, Q.dp)

length = 1.0
kn = - 2.0
ks = + 1.5
dp = 0.0001
Q = Quadrupole('Q', length, kn, ks, dp)

L = Linear('L', (dp*vector).tolist(), matrix.tolist())

state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

print(Q(state).tolist())
print(L(state).tolist())
print((Q(state) - L(state)).tolist())
[0.012284014325104016, 0.013015328226371463, 0.005866573649483532, 0.01748719571870364]
[0.012284330668424102, 0.013015814342968814, 0.005867420763768485, 0.01748647747810468]
[-3.163433200861071e-07, -4.861165973507608e-07, -8.471142849530294e-07, 7.182405989576701e-07]
[4]:
# Linear dipole

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
dp = 0.0
D = Dipole('D', length, angle, e1, e2, dp)

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)

matrix = torch.func.jacrev(lambda state, dp: D(state, data={**D.data(), **{'dp': dp}}), 0)(state, D.dp)
vector = torch.func.jacrev(lambda state, dp: D(state, data={**D.data(), **{'dp': dp}}), 1)(state, D.dp)

length = 2.0
angle = 0.05
e1 = 0.025
e2 = 0.025
dp = 0.0001
D = Dipole('D', length, angle, e1, e2, dp)

L = Linear('L', (dp*vector).tolist(), matrix.tolist())

state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

print(D(state).tolist())
print(L(state).tolist())
print((D(state) - L(state)).tolist())
[2.8326139356223135e-06, -0.005001000818821848, -0.002994615078158236, 0.001004198140800915]
[9.165104284458664e-06, -0.004994998958072864, -0.002993748697591074, 0.001004997134048961]
[-6.33249034883635e-06, -6.001860748984511e-06, -8.663805671619597e-07, -7.989932480460815e-07]

Example-14: BPM (element)

[1]:
from pathlib import Path
from os import system

import torch
from model.library.drift import Drift
from model.library.bpm import BPM
[2]:
# BPM acts as identity transformation with calibration error

state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

B = BPM('B', direction='forward')
B(state)
[2]:
tensor([ 0.0100, -0.0050, -0.0050,  0.0010], dtype=torch.float64)
[3]:
# Calibration errors can be passed as deviation variables

# qx -> (1 + xx) qx + xy qy
# qy -> yx qx + (1 + yy) qy

B.data()
[3]:
{'xx': tensor(0., dtype=torch.float64),
 'xy': tensor(0., dtype=torch.float64),
 'yx': tensor(0., dtype=torch.float64),
 'yy': tensor(0., dtype=torch.float64),
 'dp': tensor(0., dtype=torch.float64),
 'dx': tensor(0., dtype=torch.float64),
 'dy': tensor(0., dtype=torch.float64),
 'dz': tensor(0., dtype=torch.float64),
 'wx': tensor(0., dtype=torch.float64),
 'wy': tensor(0., dtype=torch.float64),
 'wz': tensor(0., dtype=torch.float64)}
[4]:
# Transform to BPM frame and back to beam frame

xx = torch.tensor(+0.05, dtype=torch.float64)
xy = torch.tensor(+0.01, dtype=torch.float64)
yx = torch.tensor(+0.05, dtype=torch.float64)
yy = torch.tensor(-0.06, dtype=torch.float64)

B = BPM('B', direction='forward')

state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
print(state)

state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})
print(state)

B.direction = 'inverse'

state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})
print(state)
tensor([ 0.0100, -0.0050, -0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0105, -0.0048, -0.0042,  0.0011], dtype=torch.float64)
tensor([ 0.0100, -0.0050, -0.0050,  0.0010], dtype=torch.float64)
[5]:
# Transform to BPM frame and back to beam frame using a pair of BPMS

xx = torch.tensor(+0.05, dtype=torch.float64)
xy = torch.tensor(+0.01, dtype=torch.float64)
yx = torch.tensor(+0.05, dtype=torch.float64)
yy = torch.tensor(-0.06, dtype=torch.float64)

BA = BPM('B', direction='forward')
BB = BPM('B', direction='inverse')

state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
print(state)

state = BA(state, data={**BA.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})
print(state)

state = BB(state, data={**BB.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})
print(state)
tensor([ 0.0100, -0.0050, -0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0105, -0.0048, -0.0042,  0.0011], dtype=torch.float64)
tensor([ 0.0100, -0.0050, -0.0050,  0.0010], dtype=torch.float64)
[6]:
# Differentiability

B = BPM('B', direction='forward')
D = Drift('D', 1.0)

xx = torch.tensor(+0.05, dtype=torch.float64)
xy = torch.tensor(+0.01, dtype=torch.float64)
yx = torch.tensor(+0.05, dtype=torch.float64)
yy = torch.tensor(-0.06, dtype=torch.float64)

error = torch.stack([xx, xy, yx, yy])

state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)

def line(state, error):
    xx, xy, yx, yy = error
    state = D(state)
    state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})
    return state

print(torch.func.jacrev(line, 1)(state, error))

def line(state, error):
    xx, xy, yx, yy = error
    state = D(state)
    B.direction = 'forward'
    state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})
    B.direction = 'inverse'
    state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}})
    return state

print(torch.func.jacrev(line, 1)(state, error))
tensor([[ 5.0000e-03, -4.0000e-03,  0.0000e+00,  0.0000e+00],
        [ 4.5880e-03, -2.4404e-04, -1.0625e-03,  5.6516e-05],
        [ 0.0000e+00,  0.0000e+00,  5.0000e-03, -4.0000e-03],
        [-4.8809e-05,  5.1249e-03,  1.1303e-05, -1.1868e-03]],
       dtype=torch.float64)
tensor([[ 0.0000e+00,  8.6736e-19,  0.0000e+00, -5.4888e-19],
        [ 0.0000e+00,  0.0000e+00,  3.3881e-20,  0.0000e+00],
        [-2.7105e-19,  0.0000e+00,  8.6736e-19,  0.0000e+00],
        [ 2.1684e-19,  3.9302e-19, -3.3881e-21,  2.1684e-19]],
       dtype=torch.float64)
[7]:
# Alignment support

xx = torch.tensor(+0.05, dtype=torch.float64)
xy = torch.tensor(+0.01, dtype=torch.float64)
yx = torch.tensor(+0.05, dtype=torch.float64)
yy = torch.tensor(-0.06, dtype=torch.float64)

dx = torch.tensor(0.05, dtype=torch.float64)
dy = torch.tensor(-0.02, dtype=torch.float64)
dz = torch.tensor(0.05, dtype=torch.float64)

wx = torch.tensor(0.005, dtype=torch.float64)
wy = torch.tensor(-0.005, dtype=torch.float64)
wz = torch.tensor(0.1, dtype=torch.float64)

error = {'dx': dx, 'dy': dy, 'dz': dz, 'wx': wx, 'wy': wy, 'wz': wz}

B = BPM('B', direction='forward')

state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}, **error})
print(state)

state = torch.tensor([0.01, -0.005, -0.005, 0.001], dtype=torch.float64)
state = B(state, data={**B.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}, **error}, alignment=True)
print(state)
tensor([ 0.0105, -0.0048, -0.0042,  0.0011], dtype=torch.float64)
tensor([ 0.0086, -0.0048, -0.0082,  0.0008], dtype=torch.float64)

Example-15: Line (element)

[1]:
# In this example line usage is illustrated
[2]:
from pathlib import Path
from os import system

from pprint import pprint

import matplotlib
from matplotlib import pyplot as plt
matplotlib.rcParams['text.usetex'] = True

import torch

from model.library.drift import Drift
from model.library.quadrupole import Quadrupole
from model.library.sextupole import Sextupole
from model.library.dipole import Dipole
from model.library.line import Line
[3]:
# Define unique elements

QF = Quadrupole('QF', 0.5, +0.25)
QD = Quadrupole('QD', 0.5, -0.20)
SF = Sextupole('SF', 0.25)
SD = Sextupole('SD', 0.25)
DR = Drift('DR', 0.25)
BM = Dipole('BM', 3.50, torch.pi/8.0)
[4]:
# Define a line
# With propagate flag, values for other flags will be propagated to all root elements (but not to lines!)

FODO = Line('FODO',
            [QF, DR, SF, DR, BM, DR, SD, DR, QD, QD, DR, SD, DR, BM, DR, SF, DR, QF],
            propagate=True,
            dp=0.0,
            exact=False,
            output=False,
            matrix = False)
[5]:
# Similar to elements, lines have to variants of call method

# 1) without deviation variables

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)
print(FODO(state))

# 2) without deviation variables (defaut dictionary)

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)
print(FODO(state, data=FODO.data()))

# In both cases, it is possible to apply alignments errors using alignment flag
# Alignment errors can be applied to elements
tensor([0., 0., 0., 0.], dtype=torch.float64)
tensor([0., 0., 0., 0.], dtype=torch.float64)
[6]:
# Values of deviation variables can be passed only to unique elements

pprint(FODO.data(alignment=False))
{'BM': {'dl': tensor(0., dtype=torch.float64),
        'dp': tensor(0., dtype=torch.float64),
        'dw': tensor(0., dtype=torch.float64),
        'e1': tensor(0., dtype=torch.float64),
        'e2': tensor(0., dtype=torch.float64),
        'kn': tensor(0., dtype=torch.float64),
        'ks': tensor(0., dtype=torch.float64),
        'mo': tensor(0., dtype=torch.float64),
        'ms': tensor(0., dtype=torch.float64)},
 'DR': {'dl': tensor(0., dtype=torch.float64),
        'dp': tensor(0., dtype=torch.float64)},
 'QD': {'dl': tensor(0., dtype=torch.float64),
        'dp': tensor(0., dtype=torch.float64),
        'kn': tensor(0., dtype=torch.float64),
        'ks': tensor(0., dtype=torch.float64)},
 'QF': {'dl': tensor(0., dtype=torch.float64),
        'dp': tensor(0., dtype=torch.float64),
        'kn': tensor(0., dtype=torch.float64),
        'ks': tensor(0., dtype=torch.float64)},
 'SD': {'dl': tensor(0., dtype=torch.float64),
        'dp': tensor(0., dtype=torch.float64),
        'ms': tensor(0., dtype=torch.float64)},
 'SF': {'dl': tensor(0., dtype=torch.float64),
        'dp': tensor(0., dtype=torch.float64),
        'ms': tensor(0., dtype=torch.float64)}}
[7]:
# Mapping over a set of initial conditions

state = 1.0E-3*torch.randn((512, 4), dtype=FODO.dtype, device=FODO.device)
print(torch.vmap(FODO)(state).shape)
torch.Size([512, 4])
[8]:
# Mapping over a set of initial conditions and parameters

state = 1.0E-3*torch.randn((512, 4), dtype=FODO.dtype, device=FODO.device)
dknqf = 1.0E-3*torch.randn(512, dtype=FODO.dtype, device=FODO.device)

def wrapper(state, dknqf):
    data = FODO.data()
    data['QF']['kn'] = dknqf
    return FODO(state, data=data)

print(torch.vmap(wrapper)(state, dknqf).shape)
torch.Size([512, 4])
[9]:
# Differentiability (state)

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)

print(torch.func.jacrev(FODO)(state))
tensor([[-0.4395, 15.4433,  0.0000,  0.0000],
        [-0.0522, -0.4395,  0.0000,  0.0000],
        [ 0.0000,  0.0000,  0.4963,  5.3596],
        [ 0.0000,  0.0000, -0.1406,  0.4963]], dtype=torch.float64)
[10]:
# Differentiability (parameter)

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)

dknqf = torch.tensor(0.0, dtype=torch.float64)

def wrapper(dknqf):
    data = FODO.data()
    data['QF']['kn'] = dknqf
    return FODO(state, data=data)

print(torch.func.jacrev(wrapper)(dknqf))
tensor([0., 0., 0., 0.], dtype=torch.float64)
[11]:
# Differentiability (composed)

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)

dknqf = torch.tensor(0.0, dtype=torch.float64)

def wrapper(dknqf):
    data = FODO.data()
    data['QF']['kn'] = dknqf
    return torch.func.jacrev(lambda state, data: FODO(state, data=data))(state, data)

print(wrapper(dknqf))
print(torch.func.jacrev(wrapper)(dknqf))
tensor([[-0.4395, 15.4433,  0.0000,  0.0000],
        [-0.0522, -0.4395,  0.0000,  0.0000],
        [ 0.0000,  0.0000,  0.4963,  5.3596],
        [ 0.0000,  0.0000, -0.1406,  0.4963]], dtype=torch.float64)
tensor([[-7.5649, -3.8172,  0.0000,  0.0000],
        [ 0.4176, -7.5649,  0.0000,  0.0000],
        [ 0.0000,  0.0000,  2.7423,  1.3262],
        [ 0.0000,  0.0000,  0.5426,  2.7423]], dtype=torch.float64)
[12]:
%%time

# Tracking

# Note, in general tracking is slow
# This is especially the case for exact integration of elements with large number of slices

qx = torch.linspace(0.0, 0.01, 16, dtype=torch.float64)
px = torch.zeros_like(qx)
qy = torch.zeros_like(qx)
py = torch.zeros_like(qx)

state = torch.stack([qx, px, qy, py]).T
orbit = []

for _ in range(2**10):
    state = torch.vmap(FODO)(state)
    orbit.append(state)

qx, px, *_ = torch.stack(orbit).swapaxes(0, -1)
CPU times: user 8.17 s, sys: 2.67 ms, total: 8.18 s
Wall time: 8.18 s
[13]:
# Plot trajectories

plt.figure(figsize=(4, 4))
plt.scatter(qx.cpu().numpy(), px.cpu().numpy(), s=1, color='black')
plt.show()
../_images/examples_model_176_0.png
[14]:
%%time

# Output can be collected at each integration step
# Note, container is overwritten at each call

FODO.ns = 0.01
FODO.output = True

state = torch.tensor([+0.01, 0.0, -0.01, 0.0], dtype=torch.float64)
orbit = []
for _ in range(16):
    state = FODO(state)
    orbit.append(FODO.container_output)
qx, _, qy, _  = torch.vstack(orbit).T

plt.figure(figsize=(8, 2))
plt.scatter(range(len(qx)), qx.cpu().numpy(), s=1, color='blue')
plt.scatter(range(len(qy)), qy.cpu().numpy(), s=1, color='red')
plt.show()
../_images/examples_model_177_0.png
CPU times: user 14 s, sys: 159 ms, total: 14.2 s
Wall time: 14 s
[15]:
%%time

# Functions can be compiled, but note that dynamo unrolls loops completely (torch 2.4)
# This leads to very long compilation times

FODO.ns = 1

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)
table = 1.0E-3*torch.randn((512, 4), dtype=FODO.dtype, device=FODO.device)

fodo = torch.compile(FODO)
fodo(state)
CPU times: user 17.8 s, sys: 5.1 s, total: 22.9 s
Wall time: 24.8 s
[15]:
tensor([0., 0., 0., 0.], dtype=torch.float64)
[16]:
%%timeit

FODO(state)
4.71 ms ± 57.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
[17]:
%%timeit

fodo(state)
458 µs ± 2.08 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
[18]:
%%timeit

_ = torch.vmap(FODO)(table)
8.5 ms ± 46.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
[19]:
%%timeit

_ = torch.vmap(fodo)(table)
8.89 ms ± 514 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
[20]:
# Note, compositional operations (map or jacobian seems to break/ignore compiled version)
[21]:
# name (property)
# This property can be used to get/set line name

FODO.name
[21]:
'FODO'
[22]:
# sequence (property)
# Contains ordered sequence of elements
# Elements can be added or removed from it

len(FODO.sequence)
[22]:
18
[23]:
# unique (property)
# name: (type, length, angle) data for all unique elements

pprint(FODO.unique)
{'BM': ('Dipole',
        tensor(3.5000, dtype=torch.float64),
        tensor(0.3927, dtype=torch.float64)),
 'DR': ('Drift',
        tensor(0.2500, dtype=torch.float64),
        tensor(0., dtype=torch.float64)),
 'QD': ('Quadrupole',
        tensor(0.5000, dtype=torch.float64),
        tensor(0., dtype=torch.float64)),
 'QF': ('Quadrupole',
        tensor(0.5000, dtype=torch.float64),
        tensor(0., dtype=torch.float64)),
 'SD': ('Sextupole',
        tensor(0.2500, dtype=torch.float64),
        tensor(0., dtype=torch.float64)),
 'SF': ('Sextupole',
        tensor(0.2500, dtype=torch.float64),
        tensor(0., dtype=torch.float64))}
[24]:
# length (property)

FODO.length
[24]:
tensor(12., dtype=torch.float64)
[25]:
# angle (property)

FODO.angle
[25]:
tensor(0.7854, dtype=torch.float64)
[26]:
# ns (property)
# This property can be used to get/set number of integration steps to unique elements

# Set value integration steps to all elements

FODO.ns = 10
print(FODO.ns)

# Set ceil(element.length/value) integration steps to each element

FODO.ns = 0.1
print(FODO.ns)

# Set by name or type

FODO.ns = (('DR', 1), ('Sextupole', 0.01))
print(FODO.ns)
{'QF': 10, 'DR': 10, 'SF': 10, 'BM': 10, 'SD': 10, 'QD': 10}
{'QF': 5, 'DR': 3, 'SF': 3, 'BM': 35, 'SD': 3, 'QD': 5}
{'QF': 5, 'DR': 1, 'SF': 25, 'BM': 35, 'SD': 25, 'QD': 5}
[27]:
# order (property)
# This property can be used to get/set integration order to unique elements

# Set value integration steps to all elements

FODO.order = 0
print(FODO.order)

# Set by name or type

FODO.order = (('BM', 1), ('Sextupole', 1))
print(FODO.order)
{'QF': 0, 'DR': 0, 'SF': 0, 'BM': 0, 'SD': 0, 'QD': 0}
{'QF': 0, 'DR': 0, 'SF': 1, 'BM': 1, 'SD': 1, 'QD': 0}
[28]:
# Nested lines

FODO = Line('FODO', [QF, DR, SF, DR, BM, DR, SD, DR, QD, QD, DR, SD, DR, BM, DR, SF, DR, QF], propagate=True, dp=0.0, exact=False, output=False, matrix = False)
RING = Line('RING', 8*[FODO], propagate=True, dp=0.0, exact=False, output=False, matrix = False)

RING.ns = 1
print(RING.ns)

RING.order = 0
print(RING.order)

# Note, sublines are not flattened

print(len(RING.sequence))
{'QF': 1, 'DR': 1, 'SF': 1, 'BM': 1, 'SD': 1, 'QD': 1}
{'QF': 0, 'DR': 0, 'SF': 0, 'BM': 0, 'SD': 0, 'QD': 0}
8
[29]:
# Tracking, mapping and differentiation is similar to flat lines

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)
RING(state)
[29]:
tensor([0., 0., 0., 0.], dtype=torch.float64)
[30]:
# Deviation table contains unique lines with unique elements (unique by names)

pprint(RING.data(alignment=False))
{'FODO': {'BM': {'dl': tensor(0., dtype=torch.float64),
                 'dp': tensor(0., dtype=torch.float64),
                 'dw': tensor(0., dtype=torch.float64),
                 'e1': tensor(0., dtype=torch.float64),
                 'e2': tensor(0., dtype=torch.float64),
                 'kn': tensor(0., dtype=torch.float64),
                 'ks': tensor(0., dtype=torch.float64),
                 'mo': tensor(0., dtype=torch.float64),
                 'ms': tensor(0., dtype=torch.float64)},
          'DR': {'dl': tensor(0., dtype=torch.float64),
                 'dp': tensor(0., dtype=torch.float64)},
          'QD': {'dl': tensor(0., dtype=torch.float64),
                 'dp': tensor(0., dtype=torch.float64),
                 'kn': tensor(0., dtype=torch.float64),
                 'ks': tensor(0., dtype=torch.float64)},
          'QF': {'dl': tensor(0., dtype=torch.float64),
                 'dp': tensor(0., dtype=torch.float64),
                 'kn': tensor(0., dtype=torch.float64),
                 'ks': tensor(0., dtype=torch.float64)},
          'SD': {'dl': tensor(0., dtype=torch.float64),
                 'dp': tensor(0., dtype=torch.float64),
                 'ms': tensor(0., dtype=torch.float64)},
          'SF': {'dl': tensor(0., dtype=torch.float64),
                 'dp': tensor(0., dtype=torch.float64),
                 'ms': tensor(0., dtype=torch.float64)}}}
[31]:
# If this can be used to pass different values of deviation variables to unique elements
# But values should match for consistent differentiation

LA = Line('LA', [QF, DR, SF, DR, BM, DR, SD, DR, QD, QD, DR, SD, DR, BM, DR, SF, DR, QF], propagate=True, dp=0.0, exact=False, output=False, matrix = False)
LB = Line('LB', [QF, DR, SF, DR, BM, DR, SD, DR, QD, QD, DR, SD, DR, BM, DR, SF, DR, QF], propagate=True, dp=0.0, exact=False, output=False, matrix = False)

RING = Line('RING', [LA, LB], propagate=True, dp=0.0, exact=False, output=False, matrix = False)

state = torch.tensor([0.01, 0.0, 0.01, 0.0], dtype=torch.float64)

kn = torch.tensor(0.01, dtype=torch.float64)

data = RING.data()

data['LA']['QF']['kn'] = kn
data['LB']['QF']['kn'] = kn

print(RING(state, data=data))

kna = torch.tensor(+0.01, dtype=torch.float64)
knb = torch.tensor(-0.01, dtype=torch.float64)

data['LA']['QF']['kn'] = kna
data['LB']['QF']['kn'] = knb

print(RING(state, data=data))
tensor([-0.0047,  0.0005, -0.0045, -0.0014], dtype=torch.float64)
tensor([-0.0055,  0.0005, -0.0048, -0.0014], dtype=torch.float64)
[32]:
# Modulation

FODO = Line('FODO', [QF, DR, SF, DR, BM, DR, SD, DR, QD, QD, DR, SD, DR, BM, DR, SF, DR, QF], propagate=True, dp=0.0, exact=False, output=False, matrix = False)

dkf = 1.0E-3*torch.randn(2**10, dtype=torch.float64)
dkd = 1.0E-3*torch.randn(2**10, dtype=torch.float64)

qx = torch.linspace(0.0, 0.01, 8, dtype=torch.float64)
px = torch.zeros_like(qx)
qy = torch.zeros_like(qx)
py = torch.zeros_like(qx)

state = torch.stack([qx, px, qy, py]).T
orbit = []

data = FODO.data()

for i in range(2**10):
    data['QF']['kn'] = dkf[i]
    data['QD']['kn'] = dkd[i]
    state = torch.vmap(lambda state: FODO(state, data=data))(state)
    orbit.append(state)

qx, px, *_ = torch.stack(orbit).swapaxes(0, -1)

plt.figure(figsize=(4, 4))
plt.scatter(qx.cpu().numpy(), px.cpu().numpy(), s=1, color='black')
plt.show()
../_images/examples_model_195_0.png
[33]:
# Matrix output and twiss parameters

from twiss import twiss
from twiss import propagate
from twiss import wolski_to_cs

QF = Quadrupole('QF', 0.5, +0.25)
QD = Quadrupole('QD', 0.5, -0.20)
SF = Sextupole('SF', 0.25)
SD = Sextupole('SD', 0.25)
DR = Drift('DR', 0.25)
BM = Dipole('BM', 3.50, torch.pi/8.0)

FODO = Line('FODO', [QF, DR, SF, DR, BM, DR, SD, DR, QD, QD, DR, SD, DR, BM, DR, SF, DR, QF], propagate=True, matrix=True)
FODO.ns = 0.01

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)

print(torch.func.jacrev(FODO)(state))
print()

out, *ms = FODO.container_matrix
for m in ms:
   out = m @ out
print(out)
print()

*_, w = twiss(out)
ws = [w]

for m in FODO.container_matrix:
    w = propagate(w, m)
    ws.append(w)

ws = torch.stack(ws)

_, bx, _, by = torch.vmap(wolski_to_cs)(ws).T

s = torch.linspace(0.0, 12.0, len(bx), dtype=torch.float64)

plt.figure(figsize=(8, 4))
plt.scatter(s.cpu().numpy(), bx.cpu().numpy(), s=1, color='blue')
plt.scatter(s.cpu().numpy(), by.cpu().numpy(), s=1, color='red')
plt.show()
tensor([[-0.4395, 15.4433,  0.0000,  0.0000],
        [-0.0522, -0.4395,  0.0000,  0.0000],
        [ 0.0000,  0.0000,  0.4963,  5.3596],
        [ 0.0000,  0.0000, -0.1406,  0.4963]], dtype=torch.float64)

tensor([[-0.4395, 15.4433,  0.0000,  0.0000],
        [-0.0522, -0.4395,  0.0000,  0.0000],
        [ 0.0000,  0.0000,  0.4963,  5.3596],
        [ 0.0000,  0.0000, -0.1406,  0.4963]], dtype=torch.float64)

../_images/examples_model_196_1.png
[34]:
# More line properties

# Define simple FODO based lattice using nested lines

DR = Drift('DR', 0.25)
BM = Dipole('BM', 3.50, torch.pi/4.0)

QF_A = Quadrupole('QF_A', 0.5, +0.20)
QD_A = Quadrupole('QD_A', 0.5, -0.19)
QF_B = Quadrupole('QF_B', 0.5, +0.20)
QD_B = Quadrupole('QD_B', 0.5, -0.19)
QF_C = Quadrupole('QF_C', 0.5, +0.20)
QD_C = Quadrupole('QD_C', 0.5, -0.19)
QF_D = Quadrupole('QF_D', 0.5, +0.20)
QD_D = Quadrupole('QD_D', 0.5, -0.19)

SF_A = Sextupole('SF_A', 0.25, 0.00)
SD_A = Sextupole('SD_A', 0.25, 0.00)
SF_B = Sextupole('SF_B', 0.25, 0.00)
SD_B = Sextupole('SD_B', 0.25, 0.00)
SF_C = Sextupole('SF_C', 0.25, 0.00)
SD_C = Sextupole('SD_C', 0.25, 0.00)
SF_D = Sextupole('SF_D', 0.25, 0.00)
SD_D = Sextupole('SD_D', 0.25, 0.00)

FODO_A = Line('FODO_A', [QF_A, DR, SF_A, DR, BM, DR, SD_A, DR, QD_A, QD_A, DR, SD_A, DR, BM, DR, SF_A, DR, QF_A], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_B = Line('FODO_B', [QF_B, DR, SF_B, DR, BM, DR, SD_B, DR, QD_B, QD_B, DR, SD_B, DR, BM, DR, SF_B, DR, QF_B], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_C = Line('FODO_C', [QF_C, DR, SF_C, DR, BM, DR, SD_C, DR, QD_C, QD_C, DR, SD_C, DR, BM, DR, SF_C, DR, QF_C], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_D = Line('FODO_D', [QF_D, DR, SF_D, DR, BM, DR, SD_D, DR, QD_D, QD_D, DR, SD_D, DR, BM, DR, SF_D, DR, QF_D], propagate=True, dp=0.0, exact=False, output=False, matrix=False)

RING = Line('RING', [FODO_A, FODO_B, FODO_C, FODO_D], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
[35]:
# Name and names

print(FODO_A.name)
print(FODO_A.names)
print()

print(RING.name)
print(RING.names)
print()
FODO_A
['QF_A', 'DR', 'SF_A', 'DR', 'BM', 'DR', 'SD_A', 'DR', 'QD_A', 'QD_A', 'DR', 'SD_A', 'DR', 'BM', 'DR', 'SF_A', 'DR', 'QF_A']

RING
['FODO_A', 'FODO_B', 'FODO_C', 'FODO_D']

[36]:
# Scan (recurcivly traverse lattice and yeild elements with matching attribute)

# All elements, since all elements have a name

print(len([*FODO_A.scan('name')]))
print(len([*RING.scan('name')]))

# All dipoles

print(len([*FODO_A.scan('angle')]))
print(len([*RING.scan('angle')]))
18
72
2
8
[37]:
# Flatten

print(len(RING))

RING.flatten()
print(len(RING))
4
72
[38]:
# First positon of an element

print(FODO_A.position('BM'))
print(FODO_A.names)
4
['QF_A', 'DR', 'SF_A', 'DR', 'BM', 'DR', 'SD_A', 'DR', 'QD_A', 'QD_A', 'DR', 'SD_A', 'DR', 'BM', 'DR', 'SF_A', 'DR', 'QF_A']
[39]:
# Get first element

print(FODO_A.start)
print(FODO_A.names)
QF_A
['QF_A', 'DR', 'SF_A', 'DR', 'BM', 'DR', 'SD_A', 'DR', 'QD_A', 'QD_A', 'DR', 'SD_A', 'DR', 'BM', 'DR', 'SF_A', 'DR', 'QF_A']
[40]:
# Set first element (rotate sequence)
# Note, first mathced occuranve is used

FODO_A.start = 'BM'
print(FODO_A.names)
['BM', 'DR', 'SD_A', 'DR', 'QD_A', 'QD_A', 'DR', 'SD_A', 'DR', 'BM', 'DR', 'SF_A', 'DR', 'QF_A', 'QF_A', 'DR', 'SF_A', 'DR']
[41]:
# Itemize (list of all elements with matching kind)

print(RING.itemize('Sextupole'))
['SF_A', 'SD_A', 'SF_B', 'SD_B', 'SF_C', 'SD_C', 'SF_D', 'SD_D']
[42]:
# Number of first level elements/lines

print(len(RING))
72
[43]:
# Get by index or name (first matched)

print(FODO_B[1])
print(FODO_B['DR'])
Drift(name="DR", length=0.25, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR", length=0.25, dp=0.0, exact=False, ns=1, order=0)
[44]:
# Set by index or name (first matched)

print(FODO_B.names)

FODO_B['SF_B'] = FODO_B['DR']
print(FODO_B.names)

FODO_B['SF_B'] = FODO_B['DR']
print(FODO_B.names)
['QF_B', 'DR', 'SF_B', 'DR', 'BM', 'DR', 'SD_B', 'DR', 'QD_B', 'QD_B', 'DR', 'SD_B', 'DR', 'BM', 'DR', 'SF_B', 'DR', 'QF_B']
['QF_B', 'DR', 'DR', 'DR', 'BM', 'DR', 'SD_B', 'DR', 'QD_B', 'QD_B', 'DR', 'SD_B', 'DR', 'BM', 'DR', 'SF_B', 'DR', 'QF_B']
['QF_B', 'DR', 'DR', 'DR', 'BM', 'DR', 'SD_B', 'DR', 'QD_B', 'QD_B', 'DR', 'SD_B', 'DR', 'BM', 'DR', 'DR', 'DR', 'QF_B']

Example-16: Inverse tracking

[1]:
import torch

from model.library.drift import Drift
from model.library.quadrupole import Quadrupole
from model.library.sextupole import Sextupole
from model.library.octupole import Octupole
from model.library.multipole import Multipole
from model.library.dipole import Dipole
from model.library.corrector import Corrector
from model.library.gradient import Gradient
from model.library.linear import Linear
from model.library.bpm import BPM
from model.library.marker import Marker
from model.library.line import Line
[2]:
# Drift

MF = Drift('MAG', length=1.0, dp=0.001, dx=0.0, dy=0.0, dz=0.0, wx=0.0, wy=0.0, wz=0.0, ns=10, exact=False)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Drift('MAG', length=1.0, dp=0.001, dx=0.0, dy=0.0, dz=0.0, wx=0.0, wy=0.0, wz=0.0, ns=10, exact=True)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Drift('MAG', length=1.0, dp=0.001, dx=0.01, dy=0.01, dz=-0.01, wx=0.001, wy=-0.001, wz=0.001, ns=10, exact=True)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Drift('MAG', length=1.0, dp=0.001, dx=0.01, dy=0.01, dz=-0.01, wx=0.001, wy=-0.001, wz=0.001, ns=10, exact=True, insertion=True)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0005, -0.0005,  0.0060,  0.0010], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([0., 0., 0., 0.], dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0005, -0.0005,  0.0060,  0.0010], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([0., 0., 0., 0.], dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0005, -0.0005,  0.0060,  0.0010], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([-8.6736e-19,  2.1684e-19,  0.0000e+00, -4.3368e-19],
       dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([-2.7015e-10, -5.0000e-04,  6.0000e-03,  1.0000e-03],
       dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([-2.6021e-18,  2.1684e-19,  8.6736e-19, -4.3368e-19],
       dtype=torch.float64)

[3]:
# Quadrupole

MF = Quadrupole('MAG', length=0.5, kn=3.0, ks=-1.0, dp=0.001, dx=0.0, dy=0.0, dz=0.0, wx=0.0, wy=0.0, wz=0.0, ns=10, exact=False)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Quadrupole('MAG', length=0.5, kn=3.0, ks=-1.0, dp=0.001, dx=0.0, dy=0.0, dz=0.0, wx=0.0, wy=0.0, wz=0.0, ns=10, exact=True)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Quadrupole('MAG', length=0.5, kn=3.0, ks=-1.0, dp=0.001, dx=0.01, dy=0.01, dz=-0.01, wx=0.001, wy=-0.001, wz=0.001, ns=10, exact=True)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Quadrupole('MAG', length=0.5, kn=3.0, ks=-1.0, dp=0.001, dx=0.01, dy=0.01, dz=-0.01, wx=0.001, wy=-0.001, wz=0.001, ns=10, exact=True, insertion=True)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([-0.0002, -0.0043,  0.0075,  0.0095], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 6.5052e-19,  0.0000e+00,  3.4694e-18, -8.6736e-19],
       dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([-0.0002, -0.0043,  0.0075,  0.0095], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([-8.6736e-19, -7.0473e-18,  6.0715e-18, -2.1684e-18],
       dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0046,  0.0135,  0.0046, -0.0031], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([-9.5410e-18,  1.1926e-18,  8.6736e-19, -7.8063e-18],
       dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0005, -0.0005,  0.0055,  0.0010], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([-7.8063e-18, -3.9031e-18, -2.6021e-18, -8.6736e-19],
       dtype=torch.float64)

[4]:
# Sextupole

MF = Sextupole('MAG', length=0.25, ms=10.0, dp=0.001, dx=0.0, dy=0.0, dz=0.0, wx=0.0, wy=0.0, wz=0.0, ns=10, exact=False)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Sextupole('MAG', length=0.25, ms=10.0, dp=0.001, dx=0.0, dy=0.0, dz=0.0, wx=0.0, wy=0.0, wz=0.0, ns=10, exact=True)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Sextupole('MAG', length=0.25, ms=10.0, dp=0.001, dx=0.01, dy=0.01, dz=-0.01, wx=0.001, wy=-0.001, wz=0.001, ns=10, exact=True)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Sextupole('MAG', length=0.25, ms=10.0, dp=0.001, dx=0.01, dy=0.01, dz=-0.01, wx=0.001, wy=-0.001, wz=0.001, ns=10, exact=True, insertion=True)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0009, -0.0005,  0.0053,  0.0010], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([0., 0., 0., 0.], dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0009, -0.0005,  0.0053,  0.0010], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([0., 0., 0., 0.], dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0009, -0.0006,  0.0053,  0.0011], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([-8.6736e-19,  1.0842e-19, -8.6736e-19, -4.3368e-19],
       dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0007, -0.0006,  0.0053,  0.0011], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([-8.6736e-19,  1.0842e-19, -8.6736e-19,  0.0000e+00],
       dtype=torch.float64)

[5]:
# Octupole

MF = Octupole('MAG', length=0.25, mo=25.0, dp=0.001, dx=0.0, dy=0.0, dz=0.0, wx=0.0, wy=0.0, wz=0.0, ns=10, exact=False)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Octupole('MAG', length=0.25, mo=25.0, dp=0.001, dx=0.0, dy=0.0, dz=0.0, wx=0.0, wy=0.0, wz=0.0, ns=10, exact=True)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Octupole('MAG', length=0.25, mo=25.0, dp=0.001, dx=0.01, dy=0.01, dz=-0.01, wx=0.001, wy=-0.001, wz=0.001, ns=10, exact=True)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Octupole('MAG', length=0.25, mo=25.0, dp=0.001, dx=0.01, dy=0.01, dz=-0.01, wx=0.001, wy=-0.001, wz=0.001, ns=10, exact=True, insertion=True)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0009, -0.0005,  0.0052,  0.0010], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([0., 0., 0., 0.], dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0009, -0.0005,  0.0052,  0.0010], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([0., 0., 0., 0.], dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0009, -0.0005,  0.0052,  0.0010], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([-8.6736e-19,  1.0842e-19,  0.0000e+00,  0.0000e+00],
       dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0008, -0.0005,  0.0052,  0.0010], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([-8.6736e-19,  0.0000e+00,  8.6736e-19,  0.0000e+00],
       dtype=torch.float64)

[6]:
# Multipole

MF = Multipole('MAG', length=0.25, kn=3.0, ks=-1.0, ms=10.0, mo=25.0, dp=0.001, dx=0.0, dy=0.0, dz=0.0, wx=0.0, wy=0.0, wz=0.0, ns=10, exact=False)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Multipole('MAG', length=0.25, kn=3.0, ks=-1.0, ms=10.0, mo=25.0, dp=0.001, dx=0.0, dy=0.0, dz=0.0, wx=0.0, wy=0.0, wz=0.0, ns=10, exact=True)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Multipole('MAG', length=0.25, kn=3.0, ks=-1.0, ms=10.0, mo=25.0, dp=0.001, dx=0.01, dy=0.01, dz=-0.01, wx=0.001, wy=-0.001, wz=0.001, ns=10, exact=True)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Multipole('MAG', length=0.25, kn=3.0, ks=-1.0, ms=10.0, mo=25.0, dp=0.001, dx=0.01, dy=0.01, dz=-0.01, wx=0.001, wy=-0.001, wz=0.001, ns=10, exact=True, insertion=True)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0006, -0.0024,  0.0057,  0.0048], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 7.5894e-18, -4.3368e-19, -2.0817e-17, -2.8189e-18],
       dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0006, -0.0024,  0.0057,  0.0048], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 7.3726e-18, -3.7947e-18, -2.0817e-17, -1.9516e-18],
       dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0019,  0.0072,  0.0051, -0.0005], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([-4.4235e-17,  3.0358e-18,  1.6480e-17, -4.1200e-18],
       dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0007, -0.0006,  0.0053,  0.0011], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([-3.7297e-17,  8.6736e-19,  9.5410e-18,  0.0000e+00],
       dtype=torch.float64)

[7]:
# Dipole

MF = Dipole('MAG', length=2.0, angle=0.1, e1=0.01, e2=0.05, kn=3.0, ks=-1.0, ms=10.0, mo=25.0, dp=0.001, dx=0.0, dy=0.0, dz=0.0, wx=0.0, wy=0.0, wz=0.0, ns=10, exact=False)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, -0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Dipole('MAG', length=2.0, angle=0.1, e1=0.01, e2=0.05, kn=3.0, ks=-1.0, ms=10.0, mo=25.0, dp=0.001, dx=0.01, dy=0.01, dz=-0.01, wx=0.001, wy=-0.001, wz=0.001, ns=10, exact=False)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, -0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Dipole('MAG', length=2.0, angle=0.1, e1=0.01, e2=0.05, kn=3.0, ks=-1.0, ms=10.0, mo=25.0, dp=0.001, dx=0.01, dy=0.01, dz=-0.01, wx=0.001, wy=-0.001, wz=0.001, ns=10, exact=False, insertion=True)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, -0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Dipole('MAG', length=2.0, angle=0.1, e1=0.01, e2=0.05, kn=1.0, ks=-1.0, ms=5.0, mo=10.0, dp=0.001, dx=0.0, dy=0.0, dz=0.0, wx=0.0, wy=0.0, wz=0.0, ns=10, exact=True)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, -0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Dipole('MAG', length=2.0, angle=0.1, e1=0.01, e2=0.05, kn=3.0, ks=-1.0, ms=10.0, mo=25.0, dp=0.001, dx=0.01, dy=0.01, dz=-0.01, wx=0.001, wy=-0.001, wz=0.001, ns=10, exact=True)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, -0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Dipole('MAG', length=2.0, angle=0.1, e1=0.01, e2=0.05, kn=3.0, ks=-1.0, ms=10.0, mo=25.0, dp=0.001, dx=0.01, dy=0.01, dz=-0.01, wx=0.001, wy=-0.001, wz=0.001, ns=10, exact=True, insertion=True)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, -0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()
tensor([ 0.0010, -0.0005,  0.0050, -0.0010], dtype=torch.float64)
tensor([-0.0116, -0.0136,  0.0722,  0.1266], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050, -0.0010], dtype=torch.float64)
tensor([ 4.0549e-17, -4.4452e-17, -1.2663e-16,  2.5045e-16],
       dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050, -0.0010], dtype=torch.float64)
tensor([ 0.0306,  0.0215, -0.0687, -0.1470], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050, -0.0010], dtype=torch.float64)
tensor([ 3.6689e-16, -8.8438e-16, -2.6489e-15,  4.6872e-15],
       dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050, -0.0010], dtype=torch.float64)
tensor([-0.0021, -0.0006,  0.0069, -0.0007], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050, -0.0010], dtype=torch.float64)
tensor([ 3.7297e-17,  3.1550e-17, -1.0408e-17,  2.1250e-17],
       dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050, -0.0010], dtype=torch.float64)
tensor([-0.0097, -0.0101,  0.0171,  0.0200], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050, -0.0010], dtype=torch.float64)
tensor([ 1.5400e-10, -2.2230e-10, -4.3635e-10,  5.0327e-10],
       dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050, -0.0010], dtype=torch.float64)
tensor([ 0.0304,  0.0215, -0.0691, -0.1473], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050, -0.0010], dtype=torch.float64)
tensor([-6.6071e-09, -4.3597e-10,  9.8294e-09, -1.9903e-08],
       dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050, -0.0010], dtype=torch.float64)
tensor([-0.0021, -0.0006,  0.0069, -0.0008], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050, -0.0010], dtype=torch.float64)
tensor([-1.6908e-12,  1.4334e-11,  2.0711e-11, -3.7469e-11],
       dtype=torch.float64)

[8]:
# Corrector

MF = Corrector('MAG', cx=0.01, cy=0.05, dp=0.001, dx=0.0, dy=0.0, dz=0.0, wx=0.0, wy=0.0, wz=0.0)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Corrector('MAG', cx=0.01, cy=0.05, dp=0.001, dx=0.01, dy=0.01, dz=-0.01, wx=0.001, wy=-0.001, wz=0.001)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([0.0010, 0.0095, 0.0050, 0.0510], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0000e+00, -4.3368e-19,  0.0000e+00,  8.6736e-19],
       dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([0.0011, 0.0095, 0.0055, 0.0510], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 8.6736e-19, -6.5052e-19,  0.0000e+00,  1.4962e-17],
       dtype=torch.float64)

[9]:
# Gradient

MF = Gradient('MAG', kn=0.1, ks=0.5, dp=0.001, dx=0.0, dy=0.0, dz=0.0, wx=0.0, wy=0.0, wz=0.0)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF = Gradient('MAG', kn=0.1, ks=0.5, dp=0.001, dx=0.01, dy=0.01, dz=-0.01, wx=0.001, wy=-0.001, wz=0.001)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([0.0010, 0.0019, 0.0050, 0.0020], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([0., 0., 0., 0.], dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0010, -0.0031,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 6.5052e-19,  1.0842e-18, -1.7347e-18, -2.1684e-19],
       dtype=torch.float64)

[10]:
# Linear

# Note, gives correct inverse for zero vector

length = 1.0
kn = - 2.0
ks = + 1.5
dp = 0.001
Q = Quadrupole('Q', length, kn, ks, dp)

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)

matrix = torch.func.jacrev(lambda state, dp: Q(state, data={**Q.data(), **{'dp': dp}}), 0)(state, Q.dp)
vector = torch.zeros_like(state)


MF = Linear('MAG', (dp*vector).tolist(), matrix.tolist(), dp=0.001, dx=0.0, dy=0.0, dz=0.0, wx=0.0, wy=0.0, wz=0.0)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()


MF = Linear('MAG', (dp*vector).tolist(), matrix.tolist(), dp=0.001, dx=0.01, dy=0.01, dz=-0.01, wx=0.001, wy=-0.001, wz=0.001)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0056,  0.0107,  0.0026, -0.0038], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([-2.8189e-18,  0.0000e+00,  5.2042e-18,  2.1684e-18],
       dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([-0.0009, -0.0046,  0.0097,  0.0083], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 8.6736e-19,  2.9273e-18, -1.7347e-18, -8.6736e-19],
       dtype=torch.float64)

[11]:
# Marker

MF = Marker('MAG', dp=0.001, dx=0.0, dy=0.0, dz=0.0, wx=0.0, wy=0.0, wz=0.0)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()

MF =  Marker('MAG', dp=0.001, dx=0.01, dy=0.01, dz=-0.01, wx=0.001, wy=-0.001, wz=0.001)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([0., 0., 0., 0.], dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 4.3368e-19,  2.1684e-19, -1.7347e-18, -4.3368e-19],
       dtype=torch.float64)

[12]:
# BPM

xx = torch.tensor(+0.05, dtype=torch.float64)
xy = torch.tensor(+0.01, dtype=torch.float64)
yx = torch.tensor(+0.05, dtype=torch.float64)
yy = torch.tensor(-0.06, dtype=torch.float64)

MF = BPM('MAG', dp=0.001, dx=0.0, dy=0.0, dz=0.0, wx=0.0, wy=0.0, wz=0.0)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data={**MF.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}}, alignment=True))
print(local := MI(local, data={**MI.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}}, alignment=True))
print(local - state)
print()

MF =  BPM('MAG', dp=0.001, dx=0.01, dy=0.01, dz=-0.01, wx=0.001, wy=-0.001, wz=0.001)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data={**MF.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}}, alignment=True))
print(local := MI(local, data={**MI.data(), **{'xx': xx, 'xy': xy, 'yx': yx, 'yy': yy}}, alignment=True))
print(local - state)
print()
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0011, -0.0005,  0.0047,  0.0011], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([-2.1684e-19,  1.0842e-19,  0.0000e+00,  0.0000e+00],
       dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0054,  0.0010], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([ 2.1684e-19,  0.0000e+00,  0.0000e+00, -2.1684e-19],
       dtype=torch.float64)

[13]:
# Line

QF = Quadrupole('QF', 1.0, +0.25)
QD = Quadrupole('QD', 1.0, -0.20)
SF = Sextupole('SF', 0.25)
SD = Sextupole('SD', 0.25)
DR = Drift('DR', 0.25)
BM = Dipole('BM', 3.50, torch.pi/8.0)

MF = Line('FODO',
          [QF, DR, SF, DR, BM, DR, SD, DR, QD, DR, SD, DR, BM, DR, SF, DR],
          propagate=True,
          dp=0.001,
          exact=False,
          output=False,
          matrix = False)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()


MF = Line('FODO',
          [QF, DR, SF, DR, BM, DR, SD, DR, QD, DR, SD, DR, BM, DR, SF, DR],
          propagate=True,
          dp=0.001,
          exact=True,
          output=False,
          matrix = False)
MI = MF.inverse()

state = torch.tensor([0.001, -0.0005, 0.005, 0.001], dtype=torch.float64)

print(local := state)
print(local := MF(local, data=MF.data(), alignment=True))
print(local := MI(local, data=MI.data(), alignment=True))
print(local - state)
print()
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([-0.0043, -0.0001,  0.0121, -0.0014], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([5.5793e-16, 1.2804e-16, 1.7347e-18, 8.6736e-19], dtype=torch.float64)

tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([-0.0380, -0.0058,  0.0121, -0.0015], dtype=torch.float64)
tensor([ 0.0010, -0.0005,  0.0050,  0.0010], dtype=torch.float64)
tensor([-9.8145e-13, -1.4483e-13,  1.8322e-12, -7.0449e-13],
       dtype=torch.float64)

Example-17: Layout

[1]:
# In this example graphical representation of (planar) layout if illustrated
# Layout data can be generated to draw layout in 1D, 2D and 3D along with other curves (given on sliced lattice)
# To avoid explicit dependencies on graphical libraries (matplotlib and plotly), results are returned as dictionaries
[2]:
# Import

import matplotlib
from matplotlib import pyplot as plt
from matplotlib.patches import Rectangle
matplotlib.rcParams['text.usetex'] = True

from plotly import graph_objects

import torch

from twiss import twiss
from twiss import propagate
from twiss import wolski_to_cs

from model.library.drift import Drift
from model.library.quadrupole import Quadrupole
from model.library.sextupole import Sextupole
from model.library.dipole import Dipole
from model.library.bpm import BPM
from model.library.line import Line

from model.command.layout import Layout
[3]:
# Define simple FODO based lattice

QF = Quadrupole('QF', 0.5, +0.20)
QD = Quadrupole('QD', 0.5, -0.19)
SF = Sextupole('SF', 0.25)
SD = Sextupole('SD', 0.25)
DR = Drift('DR', 0.25)
BM = Dipole('BM', 3.50, torch.pi/4.0)
BA = BPM('BA', direction='inverse')
BB = BPM('BB', direction='forward')

FODO = Line('FODO',
            [BA, QF, DR, SF, DR, BM, DR, SD, DR, QD, QD, DR, SD, DR, BM, DR, SF, DR, QF, BB],
            propagate=True,
            dp=0.0,
            exact=False,
            output=True,
            matrix = True)


RING = Line('RING',
            4*[FODO],
            propagate=True,
            dp=0.0,
            exact=False,
            output=True,
            matrix = True)
[4]:
# Perform lattice slicing
# Note, one step is used for zero length elements

RING.ns = 0.05
print(RING.ns)
{'BA': 1, 'QF': 10, 'DR': 5, 'SF': 5, 'BM': 70, 'SD': 5, 'QD': 10, 'BB': 1}
[5]:
# Compute transport matrices along closed orbit

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)

print(torch.func.jacrev(RING)(state))
print()

total, *ms = RING.container_matrix
for m in ms:
   total = m @ total
print(total)
print()
tensor([[ -0.3382, -17.5116,   0.0000,   0.0000],
        [  0.0506,  -0.3382,   0.0000,   0.0000],
        [  0.0000,   0.0000,  -0.2976,  -6.0422],
        [  0.0000,   0.0000,   0.1508,  -0.2976]], dtype=torch.float64)

tensor([[ -0.3382, -17.5116,   0.0000,   0.0000],
        [  0.0506,  -0.3382,   0.0000,   0.0000],
        [  0.0000,   0.0000,  -0.2976,  -6.0422],
        [  0.0000,   0.0000,   0.1508,  -0.2976]], dtype=torch.float64)

[6]:
# Compute twiss parameters at slices (entrance of each slice)

# Compute twiss from one-turn matrix at starting location

(nux, nuy), _, w = twiss(total)

# Print fractional tunes

print(f'{nux.item(), nuy.item()}')

# Propagate twiss parametes

ws = [w]
*ms, _ = RING.container_matrix
for m in ms:
    w = propagate(w, m)
    ws.append(w)
ws = torch.stack(ws)

# Convert to CS and plot beta functions

_, bx, _, by = torch.vmap(wolski_to_cs)(ws).T

plt.figure(figsize=(20, 4))
plt.plot(bx.cpu().numpy(), color='red', label=r'$\beta_x$')
plt.plot(by.cpu().numpy(), color='blue', label=r'$\beta_y$')
plt.xlabel(r'slice \#')
plt.ylabel(r'$\beta$ [m]')
plt.tight_layout()
plt.show()
(0.6950865039372625, 0.7018998427861426)
../_images/examples_model_228_1.png
[7]:
# Generate several trajectories on slice (turn off matrix computation)

RING.matrix = False

state = torch.tensor([+0.01, 0.0, -0.01, 0.0], dtype=torch.float64)

orbit = []
for _ in range(16):
    state = RING(state)
    orbit.append(RING.container_output.clone())
orbit = torch.stack(orbit)
orbit.shape
[7]:
torch.Size([16, 968, 4])
[8]:
# Set layout

layout = Layout(RING)
[9]:
# Reference orbit (lines and arcs)

# With step = None, returns globals coordinates for elements entrance faces (if the last entry is dropped)

x_entrance, y_entrance, *_ = layout.orbit(flat=False, step=None)

# Step size can be passed to generate smooth closed orbit curve
# Slicing is performed only in arc sections

x, y, *_ = layout.orbit(flat=False, step=0.01)

plt.figure(figsize=(6, 6))
plt.scatter(x_entrance.cpu().numpy(), y_entrance.cpu().numpy(), marker='o', color='black')
plt.plot(x.cpu().numpy(), y.cpu().numpy(), color='black')
plt.tight_layout()
plt.show()

# Note, rotation is anti-clockwise
../_images/examples_model_231_0.png
[10]:
# Information about each element
# (name, kind, length, angle)

names, kinds, lengths, angles = zip(*layout.line.layout())

# Element locations in global frame can be computed using orbit method

x, y, *_ = layout.orbit(step=None, flag=False, lengths=lengths, angles=angles)

plt.figure(figsize=(6, 6))
plt.errorbar(x.cpu().numpy(), y.cpu().numpy(), fmt=' ', ms=1, marker='x', color='black')
plt.tight_layout()
plt.show()
../_images/examples_model_232_0.png
[11]:
# In sliced lattice each element is splitted into one or more slices
# Information about each slice
# (name, kind, length, angle, point)
# Here point gives element entrance face location in global coordinate system

names, kinds, lengths, angles, points = layout.slicing_table()

x, y, *_ = points.T

plt.figure(figsize=(6, 6))
plt.errorbar(x.cpu().numpy(), y.cpu().numpy(), fmt=' ', ms=1, marker='x', color='black')
plt.tight_layout()
plt.show()
../_images/examples_model_233_0.png
[12]:
# 1D profile (rectangles)

_, _, lengths, *_ = layout.slicing_table()

rectangles, labels = layout.profile_1d(scale=5.0,
                                       shift=-2.5,
                                       text=True,
                                       delta=-4.0,
                                       rotation=90,
                                       exclude=['BPM'])

plt.figure(figsize=(20, 5))
plt.plot(lengths.cumsum(-1).cpu().numpy(), bx.cpu().numpy(), color='red', label=r'$\beta_x$')
plt.plot(lengths.cumsum(-1).cpu().numpy(), by.cpu().numpy(), color='blue', label=r'$\beta_y$')
plt.legend()
for rectangle in rectangles:
    plt.gca().add_patch(Rectangle(**rectangle))
for label in labels:
    plt.text(**label)
plt.ylim(-5.0, 25.0)
plt.tight_layout()
plt.show()

# Default colors and other parametes are defined in config attribute (layout class variable)
../_images/examples_model_234_0.png
[13]:
# 2D profile and data transformation

# Using angles and global positions, data on slices can be transformed
# In general, data scaling should be adjusted to present transformed data in appealing way\

# Transform data

*_, angles, points = layout.slicing_table()
angles = angles.cumsum(-1)

bx_x = torch.zeros_like(bx)
bx_y = +(0.05*bx + 0.5)
bx_z = torch.zeros_like(bx)
bx_points = torch.stack([bx_x, bx_y, bx_z]).T
bx_x, bx_y, bx_z = torch.vmap(layout.transform)(bx_points, points, angles).T

by_x = torch.zeros_like(by)
by_y = +(0.05*by + 0.5)
by_z = torch.zeros_like(by)
by_points = torch.stack([by_x, by_y, by_z]).T
by_x, by_y, by_z = torch.vmap(layout.transform)(by_points, points, angles).T

# Generate reference orbit

x, y, _ = layout.orbit(flat=False, step=0.01, start=(0, 0))

# Generate layout

blocks, labels = layout.profile_2d(start=(0, 0),
                                   delta=-0.6,
                                   linewidth=1.5,
                                   exclude=['Drift', 'BPM'])
# Plot

plt.figure(figsize=(6, 6))
plt.plot(x, y, color='black')
plt.plot(bx_x, bx_y, color='red', label=r'$\beta_x$')
plt.plot(by_x, by_y, color='blue', label=r'$\beta_y$')
for block in blocks:
    plt.errorbar(**block)
for label in labels:
    plt.text(**label)
plt.legend()
plt.xlabel(r'x [m]')
plt.ylabel(r'y [m]')
plt.tight_layout()
plt.show()
../_images/examples_model_235_0.png
[14]:
# 3D profile and data transformation

names, kinds, lengths, angles, points = layout.slicing_table()
angles = angles.cumsum(-1)

# Transform orbits

scale = 50.0
qx, _, qy, _ = orbit.swapaxes(0, 1).swapaxes(0, -1)
q_x = scale*torch.zeros_like(qx)
q_y = scale*qx
q_z = scale*qy
q_points = torch.stack([q_x, q_y, q_z]).swapaxes(0, -1)
q_x, q_y, q_z = torch.vmap(layout.transform)(q_points, points, angles).swapaxes(0, -1)

# Select data at BPMs

q_x_bpm = []
q_y_bpm = []
q_z_bpm = []
for kind, qx, qy, qz in zip(kinds, q_x.T, q_y.T, q_z.T):
    if kind == 'BPM':
        q_x_bpm.append(qx)
        q_y_bpm.append(qy)
        q_z_bpm.append(qz)
q_x_bpm = torch.stack(q_x_bpm)
q_y_bpm = torch.stack(q_y_bpm)
q_z_bpm = torch.stack(q_z_bpm)

# Generate reference orbit

x, y, z = layout.orbit(flat=False, step=0.01, start=(0, 0))

# Generate layout (can be saved as html with write_html method)

blocks = layout.profile_3d(scale=1.75)

# Plot

figure = graph_objects.Figure(
    data=[
        graph_objects.Scatter3d(
            x=x.numpy(),
            y=y.numpy(),
            z=z.numpy(),
            mode='lines',
            name='Orbit',
            line=dict(color='black',width=2.0,dash='solid'),
            opacity=0.75,
            showlegend=True
        ),
        graph_objects.Scatter3d(
            x=q_x.flatten().numpy(),
            y=q_y.flatten().numpy(),
            z=q_z.flatten().numpy(),
            mode='lines',
            name='Trajectory',
            line=dict(color='black',width=1.0,dash='solid'),
            opacity=0.50,
            showlegend=True

        ),
        graph_objects.Scatter3d(
            x=q_x_bpm.flatten().numpy(),
            y=q_y_bpm.flatten().numpy(),
            z=q_z_bpm.flatten().numpy(),
            mode='markers',
            name='Projection',
            marker=dict(color='red',size=1.5),
            opacity=0.50,
            showlegend=True
        ),
        *[graph_objects.Mesh3d(block) for block in blocks]
    ]
)
figure.update_layout(
    scene=dict(
        xaxis=dict(visible=False, range=[-20,20]),
        yaxis=dict(visible=False, range=[-20,20]),
        zaxis=dict(visible=False, range=[-5,5]),
        aspectratio=dict(x=1, y=1, z=1/4),
        annotations=[]
    ),
    margin=dict(l=0, r=0, t=0, b=0),
    legend=dict(orientation='v', x=0., y=1., xanchor='left', yanchor='top'),
    hoverlabel=dict(font_size=12, font_family="Rockwell", font_color='white'),
    legend_groupclick='toggleitem'
)
figure.show()

Example-19: Build MADX and ELEGANT style lattice

[1]:
# In this example it is demonstraited how to load MADX and ELEGANT lattice files

# Supported elements (other elemet types will be casted to drifts)

# Drift
# Quadrupole
# Sextupole
# Octupole
# Dipole
# BPM
# Marker
# Line

# Note, only basic options are translated (not translated options are ignored)
# Note, reference orbit slicing is not correct with negative bending
[2]:
from pathlib import Path

from model.command.external import load_lattice
from model.command.build import build
from model.command.layout import Layout

from plotly import graph_objects
[3]:
# Load and build simple MADX lattice

# Set lattice file path

file = Path('../../../config/initial.madx')

# Load lattice to dictionary

table = load_lattice(file)

# Build lattice

FODO = build('FODO', 'MADX', table)
print(FODO)
Marker(name="HEAD")
BPM(name="M", direction=forward)
Quadrupole(name="QD", length=0.5, kn=-0.199999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR", length=2.0, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="BM", length=1.0, angle=0.17453292519943395, e1=0.0, e2=0.0, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR", length=2.0, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="QF", length=1.0, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR", length=2.0, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="BM", length=1.0, angle=0.17453292519943395, e1=0.0, e2=0.0, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR", length=2.0, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="QD", length=0.5, kn=-0.199999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Marker(name="TAIL")
[4]:
# Load and build simple ELEGANT lattice

# Set lattice file path

file = Path('../../../config/initial.lte')

# Load lattice to dictionary

table = load_lattice(file)

# Build lattice

FODO = build('FODO', 'ELEGANT', table)
print(FODO)
Marker(name="HEAD")
BPM(name="M", direction=forward)
Quadrupole(name="QD", length=0.5, kn=-0.199999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR", length=2, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="BM", length=1.0, angle=0.17453292519943395, e1=0.0, e2=0.0, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR", length=2, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="QF", length=1.0, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR", length=2, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="BM", length=1.0, angle=0.17453292519943395, e1=0.0, e2=0.0, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR", length=2, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="QD", length=0.5, kn=-0.199999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Marker(name="TAIL")
[5]:
# Build vepp4 lattice

file = Path('../../../config/vepp4m.lte')
data = load_lattice(file)
vepp = build('VEPP4M', 'ELEGANT', data)

# Slice dipoles

vepp.ns = (('Dipole', 0.1), )

# Profile

layout = Layout(vepp)

# Generate reference orbit

x, y, z = layout.orbit(flat=False, step=None, start=(0, 0))

# Generate layout (can be saved as html with write_html method)

blocks = layout.profile_3d(scale=2.5, exclude=['Drift', 'Marker'])

# # Plot

# figure = graph_objects.Figure(
#     data=[
#         graph_objects.Scatter3d(
#             x=x.numpy(),
#             y=y.numpy(),
#             z=z.numpy(),
#             mode='lines',
#             name='Orbit',
#             line=dict(color='black',width=2.0,dash='solid'),
#             opacity=0.75,
#             showlegend=True
#         ),
#         *[graph_objects.Mesh3d(block) for block in blocks]
#     ]
# )
# figure.update_layout(
#     scene=dict(
#         xaxis=dict(visible=False, range=[-100,100]),
#         yaxis=dict(visible=False, range=[-100,100]),
#         zaxis=dict(visible=False, range=[-25,25]),
#         aspectratio=dict(x=1, y=1, z=1/4),
#         annotations=[]
#     ),
#     margin=dict(l=0, r=0, t=0, b=0),
#     legend=dict(orientation='v', x=0., y=1., xanchor='left', yanchor='top'),
#     hoverlabel=dict(font_size=12, font_family="Rockwell", font_color='white'),
#     legend_groupclick='toggleitem'
# )
# figure.show()
[6]:
# Build skif lattice

file = Path('../../../config/skif.lte')
data = load_lattice(file)
skif = build('SKIF', 'ELEGANT', data)

# Slice dipoles

skif.ns = (('Dipole', 0.1), )

# Profile

layout = Layout(skif)

# Generate reference orbit

x, y, z = layout.orbit(flat=False, step=None, start=(0, 0))

# Generate layout (can be saved as html with write_html method)

blocks = layout.profile_3d(scale=2.5, exclude=['Drift', 'Marker'])

# # Plot

# figure = graph_objects.Figure(
#     data=[
#         graph_objects.Scatter3d(
#             x=x.numpy(),
#             y=y.numpy(),
#             z=z.numpy(),
#             mode='lines',
#             name='Orbit',
#             line=dict(color='black',width=2.0,dash='solid'),
#             opacity=0.75,
#             showlegend=True
#         ),
#         *[graph_objects.Mesh3d(block) for block in blocks]
#     ]
# )
# figure.update_layout(
#     scene=dict(
#         xaxis=dict(visible=False, range=[-200,200]),
#         yaxis=dict(visible=False, range=[-200,200]),
#         zaxis=dict(visible=False, range=[-50,50]),
#         aspectratio=dict(x=1, y=1, z=1/4),
#         annotations=[]
#     ),
#     margin=dict(l=0, r=0, t=0, b=0),
#     legend=dict(orientation='v', x=0., y=1., xanchor='left', yanchor='top'),
#     hoverlabel=dict(font_size=12, font_family="Rockwell", font_color='white'),
#     legend_groupclick='toggleitem'
# )
# figure.show()

Example-20: Line serialization and deserialization (YAML)

[1]:
# In this example lattice serialization and deserialization is demonstrated
# All elements have serialize property, for all elements but lines, it returns a dictionary that can be used to construct the element
# Element(**Element.serialize)
# Line serialize property also returns a (nested) dictionary with element parameters and kinds
# This dictionary can't be used to construct the original line with constructor
# Instead, model.command.build.load_line can be used for construction (deserialization)
# This function loads and processes YAML fine (can be created with model.command.build.save_line)
[2]:
# Import

import torch

from pathlib import Path
from pprint import pprint

from model.library.drift import Drift
from model.library.quadrupole import Quadrupole
from model.library.sextupole import Sextupole
from model.library.dipole import Dipole
from model.library.bpm import BPM
from model.library.line import Line

from model.command.external import load_lattice
from model.command.build import build
from model.command.build import save_line
from model.command.build import load_line
[3]:
# Define simple FODO based lattice

QF = Quadrupole('QF', 0.5, +0.20)
QD = Quadrupole('QD', 0.5, -0.19)
SF = Sextupole('SF', 0.25)
SD = Sextupole('SD', 0.25)
DR = Drift('DR', 0.25)
BM = Dipole('BM', 3.50, torch.pi/4.0)
BA = BPM('BA', direction='inverse')
BB = BPM('BB', direction='forward')

FODO = Line('FODO',
            [BA, QF, DR, SF, DR, BM, DR, SD, DR, QD, QD, DR, SD, DR, BM, DR, SF, DR, QF, BB],
            propagate=True,
            dp=0.0,
            exact=False,
            output=True,
            matrix = True)
[4]:
# Element serialization and deserialization

pprint(DR.serialize, sort_dicts=False)
print()

print(Drift(**DR.serialize))
print()
{'name': 'DR',
 'length': 0.25,
 'dp': 0.0,
 'dx': 0.0,
 'dy': 0.0,
 'dz': 0.0,
 'wx': 0.0,
 'wy': 0.0,
 'wz': 0.0,
 'ns': 1,
 'order': 0,
 'exact': False,
 'insertion': False,
 'output': True,
 'matrix': True}

Drift(name="DR", length=0.25, dp=0.0, exact=False, ns=1, order=0)

[5]:
# Line serialization

pprint(FODO.serialize, sort_dicts=False)
{'kind': 'Line',
 'name': 'FODO',
 'sequence': [{'kind': 'BPM',
               'name': 'BA',
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'output': True,
               'matrix': True,
               'direction': 'inverse'},
              {'kind': 'Quadrupole',
               'name': 'QF',
               'length': 0.5,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': True,
               'matrix': True,
               'kn': 0.200000000000001,
               'ks': 0.0},
              {'kind': 'Drift',
               'name': 'DR',
               'length': 0.25,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': True,
               'matrix': True},
              {'kind': 'Sextupole',
               'name': 'SF',
               'length': 0.25,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': True,
               'matrix': True,
               'ms': 0.0},
              {'kind': 'Drift',
               'name': 'DR',
               'length': 0.25,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': True,
               'matrix': True},
              {'kind': 'Dipole',
               'name': 'BM',
               'length': 3.5,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': True,
               'matrix': True,
               'angle': 0.7853981633974493,
               'e1': 0.0,
               'e2': 0.0,
               'kn': 1e-15,
               'ks': 0.0,
               'ms': 0.0,
               'mo': 0.0,
               'e1_on': True,
               'e2_on': True},
              {'kind': 'Drift',
               'name': 'DR',
               'length': 0.25,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': True,
               'matrix': True},
              {'kind': 'Sextupole',
               'name': 'SD',
               'length': 0.25,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': True,
               'matrix': True,
               'ms': 0.0},
              {'kind': 'Drift',
               'name': 'DR',
               'length': 0.25,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': True,
               'matrix': True},
              {'kind': 'Quadrupole',
               'name': 'QD',
               'length': 0.5,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': True,
               'matrix': True,
               'kn': -0.189999999999999,
               'ks': 0.0},
              {'kind': 'Quadrupole',
               'name': 'QD',
               'length': 0.5,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': True,
               'matrix': True,
               'kn': -0.189999999999999,
               'ks': 0.0},
              {'kind': 'Drift',
               'name': 'DR',
               'length': 0.25,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': True,
               'matrix': True},
              {'kind': 'Sextupole',
               'name': 'SD',
               'length': 0.25,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': True,
               'matrix': True,
               'ms': 0.0},
              {'kind': 'Drift',
               'name': 'DR',
               'length': 0.25,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': True,
               'matrix': True},
              {'kind': 'Dipole',
               'name': 'BM',
               'length': 3.5,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': True,
               'matrix': True,
               'angle': 0.7853981633974493,
               'e1': 0.0,
               'e2': 0.0,
               'kn': 1e-15,
               'ks': 0.0,
               'ms': 0.0,
               'mo': 0.0,
               'e1_on': True,
               'e2_on': True},
              {'kind': 'Drift',
               'name': 'DR',
               'length': 0.25,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': True,
               'matrix': True},
              {'kind': 'Sextupole',
               'name': 'SF',
               'length': 0.25,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': True,
               'matrix': True,
               'ms': 0.0},
              {'kind': 'Drift',
               'name': 'DR',
               'length': 0.25,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': True,
               'matrix': True},
              {'kind': 'Quadrupole',
               'name': 'QF',
               'length': 0.5,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': True,
               'matrix': True,
               'kn': 0.200000000000001,
               'ks': 0.0},
              {'kind': 'BPM',
               'name': 'BB',
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'output': True,
               'matrix': True,
               'direction': 'forward'}],
 'propagate': True,
 'dp': 0.0,
 'exact': False,
 'output': True,
 'matrix': True}
[6]:
# Build lattice from ELEGANT

path = Path('./config/ic.lte')
data = load_lattice(path)
ring = build('RING', 'ELEGANT', data)

print(ring.length)
print(ring.describe)
tensor(27.4256, dtype=torch.float64)
{'BPM': 16, 'Drift': 42, 'Quadrupole': 28, 'Dipole': 8, 'Sextupole': 16}
[7]:
# Save lattice

path = Path('ring.yaml')
save_line(ring, path)
[8]:
# Load lattice

path = Path('ring.yaml')
ring = load_line(path)

print(ring.length)
print(ring.describe)
tensor(27.4256, dtype=torch.float64)
{'BPM': 16, 'Drift': 42, 'Quadrupole': 28, 'Dipole': 8, 'Sextupole': 16}

Example-21: Advanced line editing

[1]:
# In this example some advanced line addition options are explored

# Given an ELEGANT lattice the following is performed

# 1. build lattice
# 2. flatten lattice
# 3. merge drifts
# 4. split BPMs (replace BPMs with forward/inverse BPM elements)
# 5. change lattice start
# 6. split quadrupoles and insert correctors
# 7. rename inserted correctors (binding to different tensors)
# 8. generate sublines (bpm to bpm transformations)
# 9. generate layout and plot lattice

# Additionaly all line methods and properties are illustrated (also see the previous line example)
[2]:
# Import

from pprint import pprint

import torch

from pathlib import Path

from plotly import graph_objects

from model.library.line import Line
from model.library.corrector import Corrector

from model.command.external import load_lattice
from model.command.build import save_line
from model.command.build import load_line
from model.command.build import build
from model.command.layout import Layout
[3]:
# Load test ELEGANT lattice

path = Path('ic.lte')
data = load_lattice(path)
[4]:
# Build from ELEGANT style dictionary

ring:Line = build('RING', 'ELEGANT', data)
ring.propagate = True

print(f'{ring.length.item():.3f}')
print(f'{ring.angle.item():.3f}')
print(ring.describe)
27.426
6.283
{'BPM': 16, 'Drift': 42, 'Quadrupole': 28, 'Dipole': 8, 'Sextupole': 16}
[5]:
# Print full ring
# Note, lines, if any, are always flattened in printout

print(ring)
BPM(name="BPM05", direction="forward")
Drift(name="D09", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D3", length=0.19657, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D19", length=0.161715, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D38", length=0.806715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D2", length=0.19657, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D12", length=0.126715, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D42", length=0.095, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F2", length=0.18, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D28", length=0.065, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM07", direction="forward")
Drift(name="D10", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D05", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D23", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")
Drift(name="D07", length=0.111, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_3F4", length=0.08, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D30", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F4", length=0.18, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D30", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_3F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D16", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D05", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM6", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D10", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM09", direction="forward")
Drift(name="D28", length=0.065, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F1", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D14", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D1", length=0.18, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D21", length=0.36, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D21", length=0.36, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4D1", length=0.18, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D14", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F1", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D30", length=0.07, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM10", direction="forward")
Drift(name="D08", length=0.11458, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM7", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D05", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_4F4", length=0.08, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D16", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_4F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D30", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F4", length=0.18, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D30", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_4F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D11", length=0.11, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM11", direction="forward")
Drift(name="D24", length=0.045, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_4F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D05", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM8", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D10", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM12", direction="forward")
Drift(name="D28", length=0.065, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F2", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D42", length=0.095, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D12", length=0.126715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4D2", length=0.19657, kn=-2.645953025999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D39", length=0.811715, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D17", length=0.156715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4D3", length=0.19657, kn=-8.526334987999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D09", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM13", direction="forward")
Drift(name="D27", length=0.056715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F3", length=0.19657, kn=7.792169674000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D29", length=0.066715, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D06", length=0.1119, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D40", length=0.841, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D06", length=0.1119, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D29", length=0.066715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F3", length=0.19657, kn=7.792169674000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D31", length=0.071715, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM14", direction="forward")
Drift(name="D01", length=0.101715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1D3", length=0.19657, kn=-8.526334987999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D20", length=0.166715, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D36", length=0.801715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1D2", length=0.19657, kn=-2.645953025999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D15", length=0.131715, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D41", length=0.09, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F2", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D30", length=0.07, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM15", direction="forward")
Drift(name="D08", length=0.11458, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM1", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D05", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_1F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D35", length=0.0775, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM16", direction="forward")
Drift(name="D35", length=0.0775, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_1F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D30", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F4", length=0.18, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D30", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_1F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D16", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_1F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D05", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM2", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D13", length=0.12958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM17", direction="forward")
Drift(name="D26", length=0.055, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F1", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D14", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1D1", length=0.18, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D32", length=0.72, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2D1", length=0.18, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D14", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F1", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D33", length=0.075, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM01", direction="forward")
Drift(name="D05", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM3", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D05", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_2F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D16", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_2F4", length=0.08, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D30", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F4", length=0.18, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D30", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_2F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D02", length=0.101, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM02", direction="forward")
Drift(name="D25", length=0.054, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_2F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D05", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM4", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D10", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM03", direction="forward")
Drift(name="D28", length=0.065, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F2", length=0.18, kn=13.690633560000002, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D41", length=0.09, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D15", length=0.131715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2D2", length=0.19657, kn=-2.815743509999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D37", length=0.802715, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D18", length=0.159715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2D3", length=0.19657, kn=-8.487311754999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D03", length=0.106715, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM04", direction="forward")
Drift(name="D29", length=0.066715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F3", length=0.19657, kn=7.899043454000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D34", length=0.076715, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D04", length=0.1079, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D22", length=0.4205, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D22", length=0.4205, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D06", length=0.1119, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D29", length=0.066715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F3", length=0.19657, kn=7.640511954000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="D27", length=0.056715, dp=0.0, exact=False, ns=1, order=0)
[6]:
# Flatten
# Note, the above lattice is already flat (no sublines)

# Total number of first level sequence elements (with duplicates)

print(len(ring))

# Flatten sublines is any

ring.flatten()
print(len(ring))
154
154
[7]:
# Remove all elements with length != 0 and length < 1.0E-12

ring.clean((1.0E-12, None, None, None))
[8]:
# Merge adjacent drifts
# Note, this will also rename all drift elements

ring.merge(name='DR', size=3)
print(len(ring))
135
[9]:
# Print full ring

print(ring)
BPM(name="BPM05", direction="forward")
Drift(name="DR001", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D3", length=0.19657, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR002", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D2", length=0.19657, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR003", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F2", length=0.18, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR004", length=0.065, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM07", direction="forward")
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")
Drift(name="DR008", length=0.111, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_3F4", length=0.08, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR009", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F4", length=0.18, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR010", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_3F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR011", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR012", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM6", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR013", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM09", direction="forward")
Drift(name="DR014", length=0.065, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F1", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR015", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D1", length=0.18, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR016", length=0.72, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4D1", length=0.18, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR017", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F1", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR018", length=0.07, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM10", direction="forward")
Drift(name="DR019", length=0.11458, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM7", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR020", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_4F4", length=0.08, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR021", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_4F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR022", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F4", length=0.18, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR023", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_4F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR024", length=0.11, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM11", direction="forward")
Drift(name="DR025", length=0.045, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_4F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR026", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM8", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR027", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM12", direction="forward")
Drift(name="DR028", length=0.065, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F2", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR029", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4D2", length=0.19657, kn=-2.645953025999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR030", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4D3", length=0.19657, kn=-8.526334987999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR031", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM13", direction="forward")
Drift(name="DR032", length=0.056715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F3", length=0.19657, kn=7.792169674000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR033", length=1.19823, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F3", length=0.19657, kn=7.792169674000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR034", length=0.071715, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM14", direction="forward")
Drift(name="DR035", length=0.101715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1D3", length=0.19657, kn=-8.526334987999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR036", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1D2", length=0.19657, kn=-2.645953025999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR037", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F2", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR038", length=0.07, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM15", direction="forward")
Drift(name="DR039", length=0.11458, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM1", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR040", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_1F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR041", length=0.0775, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM16", direction="forward")
Drift(name="DR042", length=0.0775, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_1F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR043", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F4", length=0.18, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR044", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_1F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR045", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_1F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR046", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM2", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR047", length=0.12958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM17", direction="forward")
Drift(name="DR048", length=0.055, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F1", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR049", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1D1", length=0.18, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR050", length=0.72, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2D1", length=0.18, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR051", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F1", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR052", length=0.075, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM01", direction="forward")
Drift(name="DR053", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM3", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR054", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_2F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR055", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_2F4", length=0.08, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR056", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F4", length=0.18, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR057", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_2F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR058", length=0.101, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM02", direction="forward")
Drift(name="DR059", length=0.054, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_2F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR060", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM4", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR061", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM03", direction="forward")
Drift(name="DR062", length=0.065, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F2", length=0.18, kn=13.690633560000002, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR063", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2D2", length=0.19657, kn=-2.815743509999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR064", length=0.9624299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2D3", length=0.19657, kn=-8.487311754999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR065", length=0.106715, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM04", direction="forward")
Drift(name="DR066", length=0.066715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F3", length=0.19657, kn=7.899043454000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR067", length=1.2042300000000001, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F3", length=0.19657, kn=7.640511954000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
[10]:
# By default only one forward BPM is created for each monitor
# It transforms phase space coordinates from the beam to bpm frame (identity without deviation variables)
# To insert inverse trasformation, BPMs can be splitted with the following command
# This method can be used to split elements by type, name and/or exclude some names
# For dipoles, e1 is on for the first part and e2 is on for the last part, both are off for all middle parts
# Split of elements with non-zero length (and angle) just creates shorter elements with the same name
# Zero length elements other than BPM are not splitted

ring.split((None, ['BPM'], None, None))
print(ring)
BPM(name="BPM05", direction="forward")
BPM(name="BPM05", direction="inverse")
Drift(name="DR001", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D3", length=0.19657, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR002", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D2", length=0.19657, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR003", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F2", length=0.18, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR004", length=0.065, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM07", direction="forward")
BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")
BPM(name="BPM08", direction="inverse")
Drift(name="DR008", length=0.111, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_3F4", length=0.08, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR009", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F4", length=0.18, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR010", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_3F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR011", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR012", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM6", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR013", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM09", direction="forward")
BPM(name="BPM09", direction="inverse")
Drift(name="DR014", length=0.065, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F1", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR015", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D1", length=0.18, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR016", length=0.72, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4D1", length=0.18, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR017", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F1", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR018", length=0.07, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM10", direction="forward")
BPM(name="BPM10", direction="inverse")
Drift(name="DR019", length=0.11458, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM7", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR020", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_4F4", length=0.08, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR021", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_4F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR022", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F4", length=0.18, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR023", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_4F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR024", length=0.11, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM11", direction="forward")
BPM(name="BPM11", direction="inverse")
Drift(name="DR025", length=0.045, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_4F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR026", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM8", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR027", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM12", direction="forward")
BPM(name="BPM12", direction="inverse")
Drift(name="DR028", length=0.065, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F2", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR029", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4D2", length=0.19657, kn=-2.645953025999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR030", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4D3", length=0.19657, kn=-8.526334987999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR031", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM13", direction="forward")
BPM(name="BPM13", direction="inverse")
Drift(name="DR032", length=0.056715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F3", length=0.19657, kn=7.792169674000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR033", length=1.19823, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F3", length=0.19657, kn=7.792169674000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR034", length=0.071715, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM14", direction="forward")
BPM(name="BPM14", direction="inverse")
Drift(name="DR035", length=0.101715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1D3", length=0.19657, kn=-8.526334987999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR036", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1D2", length=0.19657, kn=-2.645953025999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR037", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F2", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR038", length=0.07, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM15", direction="forward")
BPM(name="BPM15", direction="inverse")
Drift(name="DR039", length=0.11458, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM1", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR040", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_1F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR041", length=0.0775, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM16", direction="forward")
BPM(name="BPM16", direction="inverse")
Drift(name="DR042", length=0.0775, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_1F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR043", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F4", length=0.18, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR044", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_1F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR045", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_1F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR046", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM2", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR047", length=0.12958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM17", direction="forward")
BPM(name="BPM17", direction="inverse")
Drift(name="DR048", length=0.055, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F1", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR049", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1D1", length=0.18, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR050", length=0.72, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2D1", length=0.18, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR051", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F1", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR052", length=0.075, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM01", direction="forward")
BPM(name="BPM01", direction="inverse")
Drift(name="DR053", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM3", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR054", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_2F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR055", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_2F4", length=0.08, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR056", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F4", length=0.18, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR057", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_2F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR058", length=0.101, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM02", direction="forward")
BPM(name="BPM02", direction="inverse")
Drift(name="DR059", length=0.054, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_2F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR060", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM4", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR061", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM03", direction="forward")
BPM(name="BPM03", direction="inverse")
Drift(name="DR062", length=0.065, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F2", length=0.18, kn=13.690633560000002, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR063", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2D2", length=0.19657, kn=-2.815743509999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR064", length=0.9624299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2D3", length=0.19657, kn=-8.487311754999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR065", length=0.106715, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM04", direction="forward")
BPM(name="BPM04", direction="inverse")
Drift(name="DR066", length=0.066715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F3", length=0.19657, kn=7.899043454000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR067", length=1.2042300000000001, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F3", length=0.19657, kn=7.640511954000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
[11]:
# Line start can be changed with start property or roll method

ring.roll(1)
print(ring)
BPM(name="BPM05", direction="inverse")
Drift(name="DR001", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D3", length=0.19657, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR002", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D2", length=0.19657, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR003", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F2", length=0.18, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR004", length=0.065, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM07", direction="forward")
BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")
BPM(name="BPM08", direction="inverse")
Drift(name="DR008", length=0.111, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_3F4", length=0.08, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR009", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F4", length=0.18, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR010", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_3F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR011", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR012", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM6", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR013", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM09", direction="forward")
BPM(name="BPM09", direction="inverse")
Drift(name="DR014", length=0.065, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F1", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR015", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D1", length=0.18, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR016", length=0.72, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4D1", length=0.18, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR017", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F1", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR018", length=0.07, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM10", direction="forward")
BPM(name="BPM10", direction="inverse")
Drift(name="DR019", length=0.11458, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM7", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR020", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_4F4", length=0.08, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR021", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_4F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR022", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F4", length=0.18, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR023", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_4F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR024", length=0.11, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM11", direction="forward")
BPM(name="BPM11", direction="inverse")
Drift(name="DR025", length=0.045, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_4F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR026", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM8", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR027", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM12", direction="forward")
BPM(name="BPM12", direction="inverse")
Drift(name="DR028", length=0.065, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F2", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR029", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4D2", length=0.19657, kn=-2.645953025999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR030", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4D3", length=0.19657, kn=-8.526334987999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR031", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM13", direction="forward")
BPM(name="BPM13", direction="inverse")
Drift(name="DR032", length=0.056715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F3", length=0.19657, kn=7.792169674000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR033", length=1.19823, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F3", length=0.19657, kn=7.792169674000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR034", length=0.071715, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM14", direction="forward")
BPM(name="BPM14", direction="inverse")
Drift(name="DR035", length=0.101715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1D3", length=0.19657, kn=-8.526334987999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR036", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1D2", length=0.19657, kn=-2.645953025999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR037", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F2", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR038", length=0.07, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM15", direction="forward")
BPM(name="BPM15", direction="inverse")
Drift(name="DR039", length=0.11458, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM1", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR040", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_1F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR041", length=0.0775, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM16", direction="forward")
BPM(name="BPM16", direction="inverse")
Drift(name="DR042", length=0.0775, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_1F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR043", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F4", length=0.18, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR044", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_1F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR045", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_1F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR046", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM2", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR047", length=0.12958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM17", direction="forward")
BPM(name="BPM17", direction="inverse")
Drift(name="DR048", length=0.055, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F1", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR049", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1D1", length=0.18, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR050", length=0.72, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2D1", length=0.18, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR051", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F1", length=0.18, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR052", length=0.075, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM01", direction="forward")
BPM(name="BPM01", direction="inverse")
Drift(name="DR053", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM3", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR054", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_2F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR055", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_2F4", length=0.08, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR056", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F4", length=0.18, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR057", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_2F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR058", length=0.101, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM02", direction="forward")
BPM(name="BPM02", direction="inverse")
Drift(name="DR059", length=0.054, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_2F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR060", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM4", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR061", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM03", direction="forward")
BPM(name="BPM03", direction="inverse")
Drift(name="DR062", length=0.065, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F2", length=0.18, kn=13.690633560000002, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR063", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2D2", length=0.19657, kn=-2.815743509999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR064", length=0.9624299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2D3", length=0.19657, kn=-8.487311754999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR065", length=0.106715, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM04", direction="forward")
BPM(name="BPM04", direction="inverse")
Drift(name="DR066", length=0.066715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F3", length=0.19657, kn=7.899043454000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR067", length=1.2042300000000001, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F3", length=0.19657, kn=7.640511954000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM05", direction="forward")
[12]:
# Insert correctors at quadrupole centers
# This inserts elements with identical names
# Element name is like a hash that is used to bind tensors to deviation variables
# Thus, to bind different tensors, elements should have different names
# For example, we would like splitted BPMs and pelement part to have identical names

insection = Corrector(name='CXY')
ring.split((1 + 1, ['Quadrupole'], None, None), paste=[insection])
print(ring)
BPM(name="BPM05", direction="inverse")
Drift(name="DR001", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D3", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q3D3", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR002", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D2", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q3D2", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR003", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F2", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q3F2", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR004", length=0.065, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM07", direction="forward")
BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")
BPM(name="BPM08", direction="inverse")
Drift(name="DR008", length=0.111, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_3F4", length=0.08, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR009", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F4", length=0.09, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q3F4", length=0.09, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR010", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_3F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR011", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR012", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM6", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR013", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM09", direction="forward")
BPM(name="BPM09", direction="inverse")
Drift(name="DR014", length=0.065, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F1", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q3F1", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR015", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D1", length=0.09, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q3D1", length=0.09, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR016", length=0.72, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4D1", length=0.09, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q4D1", length=0.09, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR017", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F1", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q4F1", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR018", length=0.07, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM10", direction="forward")
BPM(name="BPM10", direction="inverse")
Drift(name="DR019", length=0.11458, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM7", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR020", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_4F4", length=0.08, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR021", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_4F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR022", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F4", length=0.09, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q4F4", length=0.09, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR023", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_4F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR024", length=0.11, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM11", direction="forward")
BPM(name="BPM11", direction="inverse")
Drift(name="DR025", length=0.045, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_4F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR026", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM8", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR027", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM12", direction="forward")
BPM(name="BPM12", direction="inverse")
Drift(name="DR028", length=0.065, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F2", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q4F2", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR029", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4D2", length=0.098285, kn=-2.645953025999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q4D2", length=0.098285, kn=-2.645953025999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR030", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4D3", length=0.098285, kn=-8.526334987999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q4D3", length=0.098285, kn=-8.526334987999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR031", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM13", direction="forward")
BPM(name="BPM13", direction="inverse")
Drift(name="DR032", length=0.056715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F3", length=0.098285, kn=7.792169674000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q4F3", length=0.098285, kn=7.792169674000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR033", length=1.19823, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F3", length=0.098285, kn=7.792169674000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q1F3", length=0.098285, kn=7.792169674000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR034", length=0.071715, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM14", direction="forward")
BPM(name="BPM14", direction="inverse")
Drift(name="DR035", length=0.101715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1D3", length=0.098285, kn=-8.526334987999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q1D3", length=0.098285, kn=-8.526334987999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR036", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1D2", length=0.098285, kn=-2.645953025999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q1D2", length=0.098285, kn=-2.645953025999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR037", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F2", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q1F2", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR038", length=0.07, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM15", direction="forward")
BPM(name="BPM15", direction="inverse")
Drift(name="DR039", length=0.11458, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM1", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR040", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_1F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR041", length=0.0775, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM16", direction="forward")
BPM(name="BPM16", direction="inverse")
Drift(name="DR042", length=0.0775, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_1F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR043", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F4", length=0.09, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q1F4", length=0.09, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR044", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_1F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR045", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_1F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR046", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM2", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR047", length=0.12958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM17", direction="forward")
BPM(name="BPM17", direction="inverse")
Drift(name="DR048", length=0.055, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F1", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q1F1", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR049", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1D1", length=0.09, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q1D1", length=0.09, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR050", length=0.72, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2D1", length=0.09, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q2D1", length=0.09, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR051", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F1", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q2F1", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR052", length=0.075, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM01", direction="forward")
BPM(name="BPM01", direction="inverse")
Drift(name="DR053", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM3", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR054", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_2F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR055", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_2F4", length=0.08, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR056", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F4", length=0.09, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q2F4", length=0.09, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR057", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_2F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR058", length=0.101, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM02", direction="forward")
BPM(name="BPM02", direction="inverse")
Drift(name="DR059", length=0.054, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_2F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR060", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM4", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR061", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM03", direction="forward")
BPM(name="BPM03", direction="inverse")
Drift(name="DR062", length=0.065, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F2", length=0.09, kn=13.690633560000002, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q2F2", length=0.09, kn=13.690633560000002, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR063", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2D2", length=0.098285, kn=-2.815743509999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q2D2", length=0.098285, kn=-2.815743509999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR064", length=0.9624299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2D3", length=0.098285, kn=-8.487311754999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q2D3", length=0.098285, kn=-8.487311754999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR065", length=0.106715, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM04", direction="forward")
BPM(name="BPM04", direction="inverse")
Drift(name="DR066", length=0.066715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F3", length=0.098285, kn=7.899043454000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q2F3", length=0.098285, kn=7.899043454000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR067", length=1.2042300000000001, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F3", length=0.098285, kn=7.640511954000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q3F3", length=0.098285, kn=7.640511954000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM05", direction="forward")
[13]:
# Individual elements can be renamed with rename method (renames the first occurrance of the element in the sequence)
# To rename by kind mangle method can be used
# Note, a list of element names to skip can be also passed to it

ring.mangle('Corrector')
print(ring)
BPM(name="BPM05", direction="inverse")
Drift(name="DR001", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D3", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_001", cx=0.0, cy=0.0)
Quadrupole(name="Q3D3", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR002", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D2", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_002", cx=0.0, cy=0.0)
Quadrupole(name="Q3D2", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR003", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F2", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_003", cx=0.0, cy=0.0)
Quadrupole(name="Q3F2", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR004", length=0.065, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM07", direction="forward")
BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")
BPM(name="BPM08", direction="inverse")
Drift(name="DR008", length=0.111, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_3F4", length=0.08, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR009", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F4", length=0.09, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_004", cx=0.0, cy=0.0)
Quadrupole(name="Q3F4", length=0.09, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR010", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_3F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR011", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR012", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM6", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR013", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM09", direction="forward")
BPM(name="BPM09", direction="inverse")
Drift(name="DR014", length=0.065, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F1", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_005", cx=0.0, cy=0.0)
Quadrupole(name="Q3F1", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR015", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D1", length=0.09, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_006", cx=0.0, cy=0.0)
Quadrupole(name="Q3D1", length=0.09, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR016", length=0.72, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4D1", length=0.09, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_007", cx=0.0, cy=0.0)
Quadrupole(name="Q4D1", length=0.09, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR017", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F1", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_008", cx=0.0, cy=0.0)
Quadrupole(name="Q4F1", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR018", length=0.07, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM10", direction="forward")
BPM(name="BPM10", direction="inverse")
Drift(name="DR019", length=0.11458, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM7", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR020", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_4F4", length=0.08, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR021", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_4F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR022", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F4", length=0.09, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_009", cx=0.0, cy=0.0)
Quadrupole(name="Q4F4", length=0.09, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR023", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_4F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR024", length=0.11, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM11", direction="forward")
BPM(name="BPM11", direction="inverse")
Drift(name="DR025", length=0.045, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_4F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR026", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM8", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR027", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM12", direction="forward")
BPM(name="BPM12", direction="inverse")
Drift(name="DR028", length=0.065, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F2", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_010", cx=0.0, cy=0.0)
Quadrupole(name="Q4F2", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR029", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4D2", length=0.098285, kn=-2.645953025999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_011", cx=0.0, cy=0.0)
Quadrupole(name="Q4D2", length=0.098285, kn=-2.645953025999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR030", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4D3", length=0.098285, kn=-8.526334987999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_012", cx=0.0, cy=0.0)
Quadrupole(name="Q4D3", length=0.098285, kn=-8.526334987999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR031", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM13", direction="forward")
BPM(name="BPM13", direction="inverse")
Drift(name="DR032", length=0.056715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q4F3", length=0.098285, kn=7.792169674000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_013", cx=0.0, cy=0.0)
Quadrupole(name="Q4F3", length=0.098285, kn=7.792169674000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR033", length=1.19823, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F3", length=0.098285, kn=7.792169674000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_014", cx=0.0, cy=0.0)
Quadrupole(name="Q1F3", length=0.098285, kn=7.792169674000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR034", length=0.071715, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM14", direction="forward")
BPM(name="BPM14", direction="inverse")
Drift(name="DR035", length=0.101715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1D3", length=0.098285, kn=-8.526334987999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_015", cx=0.0, cy=0.0)
Quadrupole(name="Q1D3", length=0.098285, kn=-8.526334987999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR036", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1D2", length=0.098285, kn=-2.645953025999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_016", cx=0.0, cy=0.0)
Quadrupole(name="Q1D2", length=0.098285, kn=-2.645953025999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR037", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F2", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_017", cx=0.0, cy=0.0)
Quadrupole(name="Q1F2", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR038", length=0.07, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM15", direction="forward")
BPM(name="BPM15", direction="inverse")
Drift(name="DR039", length=0.11458, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM1", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR040", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_1F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR041", length=0.0775, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM16", direction="forward")
BPM(name="BPM16", direction="inverse")
Drift(name="DR042", length=0.0775, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_1F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR043", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F4", length=0.09, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_018", cx=0.0, cy=0.0)
Quadrupole(name="Q1F4", length=0.09, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR044", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_1F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR045", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_1F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR046", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM2", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR047", length=0.12958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM17", direction="forward")
BPM(name="BPM17", direction="inverse")
Drift(name="DR048", length=0.055, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1F1", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_019", cx=0.0, cy=0.0)
Quadrupole(name="Q1F1", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR049", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q1D1", length=0.09, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_020", cx=0.0, cy=0.0)
Quadrupole(name="Q1D1", length=0.09, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR050", length=0.72, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2D1", length=0.09, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_021", cx=0.0, cy=0.0)
Quadrupole(name="Q2D1", length=0.09, kn=-6.763721520999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR051", length=0.12, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F1", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_022", cx=0.0, cy=0.0)
Quadrupole(name="Q2F1", length=0.09, kn=13.562217330000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR052", length=0.075, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM01", direction="forward")
BPM(name="BPM01", direction="inverse")
Drift(name="DR053", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM3", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR054", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_2F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR055", length=0.155, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX1_2F4", length=0.08, ms=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR056", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F4", length=0.09, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_023", cx=0.0, cy=0.0)
Quadrupole(name="Q2F4", length=0.09, kn=12.030967120000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR057", length=0.07, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SX2_2F4", length=0.08, ms=206.44984, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR058", length=0.101, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM02", direction="forward")
BPM(name="BPM02", direction="inverse")
Drift(name="DR059", length=0.054, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY2_2F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR060", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM4", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR061", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM03", direction="forward")
BPM(name="BPM03", direction="inverse")
Drift(name="DR062", length=0.065, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F2", length=0.09, kn=13.690633560000002, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_024", cx=0.0, cy=0.0)
Quadrupole(name="Q2F2", length=0.09, kn=13.690633560000002, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR063", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2D2", length=0.098285, kn=-2.815743509999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_025", cx=0.0, cy=0.0)
Quadrupole(name="Q2D2", length=0.098285, kn=-2.815743509999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR064", length=0.9624299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2D3", length=0.098285, kn=-8.487311754999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_026", cx=0.0, cy=0.0)
Quadrupole(name="Q2D3", length=0.098285, kn=-8.487311754999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR065", length=0.106715, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM04", direction="forward")
BPM(name="BPM04", direction="inverse")
Drift(name="DR066", length=0.066715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F3", length=0.098285, kn=7.899043454000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_027", cx=0.0, cy=0.0)
Quadrupole(name="Q2F3", length=0.098285, kn=7.899043454000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR067", length=1.2042300000000001, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F3", length=0.098285, kn=7.640511954000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_028", cx=0.0, cy=0.0)
Quadrupole(name="Q3F3", length=0.098285, kn=7.640511954000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM05", direction="forward")
[14]:
# Create lines between BPMs
# Note, printout is always flat

ring.splice()

print(len(ring))

16
[15]:
# Print the first subline

line, *_ = ring
print(line)
BPM(name="BPM05", direction="inverse")
Drift(name="DR001", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D3", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_001", cx=0.0, cy=0.0)
Quadrupole(name="Q3D3", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR002", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D2", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_002", cx=0.0, cy=0.0)
Quadrupole(name="Q3D2", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR003", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F2", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_003", cx=0.0, cy=0.0)
Quadrupole(name="Q3F2", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR004", length=0.065, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM07", direction="forward")
[16]:
# Print the second subline

_, line, *_ = ring
print(line)
BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")
[17]:
# Print the last subline

*_, line = ring
print(line)
BPM(name="BPM04", direction="inverse")
Drift(name="DR066", length=0.066715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q2F3", length=0.098285, kn=7.899043454000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_027", cx=0.0, cy=0.0)
Quadrupole(name="Q2F3", length=0.098285, kn=7.899043454000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR067", length=1.2042300000000001, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F3", length=0.098285, kn=7.640511954000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_028", cx=0.0, cy=0.0)
Quadrupole(name="Q3F3", length=0.098285, kn=7.640511954000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM05", direction="forward")
[18]:
# save/load lattice

path = Path('ring.yaml')
save_line(ring, path)
ring:Line = load_line(path)

print(f'{ring.length.item():.3f}')
print(f'{ring.angle.item():.3f}')
print(ring.describe)
27.369
6.283
{'BPM': 16, 'Drift': 67, 'Quadrupole': 28, 'Corrector': 28, 'Dipole': 8, 'Sextupole': 16}
[19]:
# Here are all the steps together without save/load steps

path = Path('ic.lte')                                              # -- set path with ELEGANT lattice file
data = load_lattice(path)                                          # -- load ELEGANT file as dictionary
ring:Line = build('RING', 'ELEGANT', data)                         # -- build ELEGANT style dictionary
ring.propagate = True                                              # -- set flag to propagate options (slicing, integration order, exact, output, ...)
ring.flatten()                                                     # -- flatten all sublines
ring.clean((1.0E-12, None, None, None))                            # -- remove non-zero length elements with length smaller than given
ring.merge()                                                       # -- merge drifts
ring.split((None, ['BPM'], None, None))                            # -- split BPMs
ring.roll(1)                                                       # -- change line start
insection = Corrector('CXY')                                       # -- define insertion element
ring.split((1 + 1, ['Quadrupole'], None, None), paste=[insection]) # -- split quadrupoles and insert corrector between parts
ring.mangle('Corrector')                                           # -- rename all correctors (unique names)
ring.splice()                                                      # -- generate lines between BPMs
[20]:
# Plot 3D profile

# Slice dipoles

ring.ns = (('Dipole', 0.05), )

# Set layout

layout = Layout(ring)

# Generate orbit

x, y, z = layout.orbit(flat=False, step=0.05, start=(0, 0))

# Generate layout blocks

blocks = layout.profile_3d(scale=2.5, exclude=['Drift'])

# Plot

figure = graph_objects.Figure(
    data=[
        graph_objects.Scatter3d(
            x=x.numpy(),
            y=y.numpy(),
            z=z.numpy(),
            mode='lines',
            name='Orbit',
            line=dict(color='black',width=2.0,dash='solid'),
            opacity=0.75,
            showlegend=True
        ),
        *[graph_objects.Mesh3d(block) for block in blocks]
    ]
)
figure.update_layout(
    scene=dict(
        xaxis=dict(visible=False, range=[-10,10]),
        yaxis=dict(visible=False, range=[-10,10]),
        zaxis=dict(visible=False, range=[-2,2]),
        aspectratio=dict(x=1, y=1, z=1/5),
        annotations=[]
    ),
    margin=dict(l=0, r=0, t=0, b=0),
    legend=dict(orientation='v', x=0., y=1., xanchor='left', yanchor='top'),
    hoverlabel=dict(font_size=12, font_family="Rockwell", font_color='white'),
    legend_groupclick='toggleitem'
)
figure.show()
[21]:
# In the rest of the example we go through all line methods and properties one by one
# Also see element properties that are not reimplemented (flag, clone)

insection = Corrector('CXY')

path = Path('ic.lte')
data = load_lattice(path)
ring:Line = build('RING', 'ELEGANT', data)
ring.propagate = True
ring.flatten()
ring.clean((1.0E-12, None, None, None))
ring.merge()
ring.split((None, ['BPM'], None, None))
ring.roll(1)
ring.split((1 + 1, ['Quadrupole'], None, None), paste=[insection])
ring.mangle('Corrector')
ring.splice()
[22]:
# serialize (property)

print([*map(len, ring)])

_, line, *_ = ring
pprint(line.serialize, sort_dicts=False)
[15, 7, 15, 19, 15, 7, 15, 11, 15, 7, 15, 19, 15, 7, 15, 10]
{'kind': 'Line',
 'name': 'BPM07_BPM08',
 'sequence': [{'kind': 'BPM',
               'name': 'BPM07',
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'output': False,
               'matrix': False,
               'direction': 'inverse'},
              {'kind': 'Drift',
               'name': 'DR005',
               'length': 0.11958,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': False,
               'matrix': False},
              {'kind': 'Dipole',
               'name': 'RM5',
               'length': 0.87284,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': False,
               'matrix': False,
               'angle': 0.785398163400001,
               'e1': 0.0,
               'e2': 0.0,
               'kn': -2.379107171999999,
               'ks': 0.0,
               'ms': 0.0,
               'mo': 0.0,
               'e1_on': True,
               'e2_on': True},
              {'kind': 'Drift',
               'name': 'DR006',
               'length': 0.10958,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': False,
               'matrix': False},
              {'kind': 'Sextupole',
               'name': 'SY1_3F4',
               'length': 0.08,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': False,
               'matrix': False,
               'ms': -277.23165},
              {'kind': 'Drift',
               'name': 'DR007',
               'length': 0.044,
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'ns': 1,
               'order': 0,
               'exact': False,
               'insertion': False,
               'output': False,
               'matrix': False},
              {'kind': 'BPM',
               'name': 'BPM08',
               'dp': 0.0,
               'dx': 0.0,
               'dy': 0.0,
               'dz': 0.0,
               'wx': 0.0,
               'wy': 0.0,
               'wz': 0.0,
               'output': False,
               'matrix': False,
               'direction': 'forward'}],
 'propagate': False,
 'dp': 0.0,
 'exact': False,
 'output': False,
 'matrix': False}
[23]:
# inverse

forward, *_ = ring.clone()
forward.propagate = True
forward.dp = 0.001

inverse = forward.inverse()

state = torch.tensor([0.001, 0.0005, -0.005, 0.0001], dtype=torch.float64)
local = forward(state)
local = inverse(local)

assert torch.allclose(local, state, rtol=1.0E-15, atol=1.0E-15)
[24]:
# data (default deviation data)

_, line, *_ = ring.clone()
pprint(line.data(alignment=False), sort_dicts=False)
{'BPM07': {'xx': tensor(0., dtype=torch.float64),
           'xy': tensor(0., dtype=torch.float64),
           'yx': tensor(0., dtype=torch.float64),
           'yy': tensor(0., dtype=torch.float64),
           'dp': tensor(0., dtype=torch.float64)},
 'DR005': {'dp': tensor(0., dtype=torch.float64),
           'dl': tensor(0., dtype=torch.float64)},
 'RM5': {'dw': tensor(0., dtype=torch.float64),
         'e1': tensor(0., dtype=torch.float64),
         'e2': tensor(0., dtype=torch.float64),
         'kn': tensor(0., dtype=torch.float64),
         'ks': tensor(0., dtype=torch.float64),
         'ms': tensor(0., dtype=torch.float64),
         'mo': tensor(0., dtype=torch.float64),
         'dp': tensor(0., dtype=torch.float64),
         'dl': tensor(0., dtype=torch.float64)},
 'DR006': {'dp': tensor(0., dtype=torch.float64),
           'dl': tensor(0., dtype=torch.float64)},
 'SY1_3F4': {'ms': tensor(0., dtype=torch.float64),
             'dp': tensor(0., dtype=torch.float64),
             'dl': tensor(0., dtype=torch.float64)},
 'DR007': {'dp': tensor(0., dtype=torch.float64),
           'dl': tensor(0., dtype=torch.float64)},
 'BPM08': {'xx': tensor(0., dtype=torch.float64),
           'xy': tensor(0., dtype=torch.float64),
           'yx': tensor(0., dtype=torch.float64),
           'yy': tensor(0., dtype=torch.float64),
           'dp': tensor(0., dtype=torch.float64)}}
[25]:
# scan (generator to get elements at all levels that have given attribute)

pprint([*ring.scan('angle')])
[Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0),
 Dipole(name="RM6", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0),
 Dipole(name="RM7", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0),
 Dipole(name="RM8", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0),
 Dipole(name="RM1", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0),
 Dipole(name="RM2", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0),
 Dipole(name="RM3", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0),
 Dipole(name="RM4", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)]
[26]:
# select (static method to filter elements)

elements = [*ring.scan('angle')]

pprint(ring.select(elements, kinds=None, names=["RM1", "RM2"]))
[Dipole(name="RM1", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0),
 Dipole(name="RM2", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)]
[27]:
# get (combination of scan and select to get attribute value)

line, *_ = ring.clone()
pprint(line.get('length', kinds=['Quadrupole']))
[('Q3D3', tensor(0.0983, dtype=torch.float64)),
 ('Q3D3', tensor(0.0983, dtype=torch.float64)),
 ('Q3D2', tensor(0.0983, dtype=torch.float64)),
 ('Q3D2', tensor(0.0983, dtype=torch.float64)),
 ('Q3F2', tensor(0.0900, dtype=torch.float64)),
 ('Q3F2', tensor(0.0900, dtype=torch.float64))]
[28]:
# set (combination of scan and select to set attribute value)
# Note, this method sets attribute to the first matched element (by name)

line, *_ = ring.clone()
line.set('length', 0.0, kinds=['Quadrupole'])
pprint(line.get('length', kinds=['Quadrupole']))
[('Q3D3', tensor(0.0983, dtype=torch.float64)),
 ('Q3D3', tensor(0., dtype=torch.float64)),
 ('Q3D2', tensor(0.0983, dtype=torch.float64)),
 ('Q3D2', tensor(0., dtype=torch.float64)),
 ('Q3F2', tensor(0.0900, dtype=torch.float64)),
 ('Q3F2', tensor(0., dtype=torch.float64))]
[29]:
# name

print(ring.name)
RING
[30]:
# sequence (ordered elements)

line, *_ = ring.clone()
pprint(line.sequence)
[BPM(name="BPM05", direction="inverse"),
 Drift(name="DR001", length=0.116715, dp=0.0, exact=False, ns=1, order=0),
 Quadrupole(name="Q3D3", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0),
 Corrector(name="CXY_001", cx=0.0, cy=0.0),
 Quadrupole(name="Q3D3", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0),
 Drift(name="DR002", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0),
 Quadrupole(name="Q3D2", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0),
 Corrector(name="CXY_002", cx=0.0, cy=0.0),
 Quadrupole(name="Q3D2", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0),
 Drift(name="DR003", length=0.221715, dp=0.0, exact=False, ns=1, order=0),
 Quadrupole(name="Q3F2", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0),
 Corrector(name="CXY_003", cx=0.0, cy=0.0),
 Quadrupole(name="Q3F2", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0),
 Drift(name="DR004", length=0.065, dp=0.0, exact=False, ns=1, order=0),
 BPM(name="BPM07", direction="forward")]
[31]:
# flatten

line = ring.clone()
print(len(line))
line.flatten()
print(len(line))
16
207
[32]:
# Rename (get/set)

line = ring.clone()
print(line[0].name)
line[0].name = 'LINE'
print(line[0].name)
BPM05_BPM07
LINE
[33]:
# append
# Note, since line sequence is basicaly a (nested) list, all list methods can be used on it
# Not all methods are implemented

line, *_ = ring.clone()
line.append(Corrector('CXY'))
print(line)
BPM(name="BPM05", direction="inverse")
Drift(name="DR001", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D3", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_001", cx=0.0, cy=0.0)
Quadrupole(name="Q3D3", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR002", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D2", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_002", cx=0.0, cy=0.0)
Quadrupole(name="Q3D2", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR003", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F2", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_003", cx=0.0, cy=0.0)
Quadrupole(name="Q3F2", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR004", length=0.065, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM07", direction="forward")
Corrector(name="CXY", cx=0.0, cy=0.0)
[34]:
# extend

line, *_ = ring.clone()
line.extend(Line('CXY', sequence=[Corrector('CXY')]))
print(line)
BPM(name="BPM05", direction="inverse")
Drift(name="DR001", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D3", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_001", cx=0.0, cy=0.0)
Quadrupole(name="Q3D3", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR002", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D2", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_002", cx=0.0, cy=0.0)
Quadrupole(name="Q3D2", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR003", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F2", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_003", cx=0.0, cy=0.0)
Quadrupole(name="Q3F2", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR004", length=0.065, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM07", direction="forward")
Corrector(name="CXY", cx=0.0, cy=0.0)
[35]:
# insert (after element with given name)

line, *_ = ring.clone()
line.insert(Corrector('CXY'), 'DR001')
print(line)
BPM(name="BPM05", direction="inverse")
Drift(name="DR001", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Quadrupole(name="Q3D3", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_001", cx=0.0, cy=0.0)
Quadrupole(name="Q3D3", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR002", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D2", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_002", cx=0.0, cy=0.0)
Quadrupole(name="Q3D2", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR003", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F2", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_003", cx=0.0, cy=0.0)
Quadrupole(name="Q3F2", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR004", length=0.065, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM07", direction="forward")
[36]:
# remove (first occurrance with matching name)

line, *_ = ring.clone()
line.remove('Q3D3')
print(line)
BPM(name="BPM05", direction="inverse")
Drift(name="DR001", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_001", cx=0.0, cy=0.0)
Quadrupole(name="Q3D3", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR002", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D2", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_002", cx=0.0, cy=0.0)
Quadrupole(name="Q3D2", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR003", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F2", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_003", cx=0.0, cy=0.0)
Quadrupole(name="Q3F2", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR004", length=0.065, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM07", direction="forward")
[37]:
# replace (first occurrance with matching name)

line, *_ = ring.clone()
line.replace('Q3D3', Corrector('CXY'))
print(line)
BPM(name="BPM05", direction="inverse")
Drift(name="DR001", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Corrector(name="CXY_001", cx=0.0, cy=0.0)
Quadrupole(name="Q3D3", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR002", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D2", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_002", cx=0.0, cy=0.0)
Quadrupole(name="Q3D2", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR003", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F2", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_003", cx=0.0, cy=0.0)
Quadrupole(name="Q3F2", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR004", length=0.065, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM07", direction="forward")
[38]:
# names (all first level names)

line, *_ = ring.clone()
pprint(line.names)
['BPM05',
 'DR001',
 'Q3D3',
 'CXY_001',
 'Q3D3',
 'DR002',
 'Q3D2',
 'CXY_002',
 'Q3D2',
 'DR003',
 'Q3F2',
 'CXY_003',
 'Q3F2',
 'DR004',
 'BPM07']
[39]:
# position (first matching position at the first level)

line, *_ = ring.clone()
line.position('Q3D3')
[39]:
2
[40]:
# positions (all matching position at the first level)

line, *_ = ring.clone()
line.positions('Q3D3')
[40]:
[2, 4]
[41]:
# start (get/set start by name, first match)

_, line, *_ = ring.clone()

print(line.start)
print(line)
print()

line.start = 'RM5'

print(line.start)
print(line)
print()
BPM07
BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")

RM5
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")
BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)

[42]:
# roll (first level sequence)

_, line, *_ = ring.clone()

print(line.start)
print(line)
print()

line.roll(line.position('RM5'))

print(line.start)
print(line)
print()
BPM07
BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")

RM5
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")
BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)

[43]:
# unique (unique elements at all levels)

_, line, *_ = ring.clone()
pprint(line.unique)
{'BPM07': ('BPM',
           tensor(0., dtype=torch.float64),
           tensor(0., dtype=torch.float64)),
 'BPM08': ('BPM',
           tensor(0., dtype=torch.float64),
           tensor(0., dtype=torch.float64)),
 'DR005': ('Drift',
           tensor(0.1196, dtype=torch.float64),
           tensor(0., dtype=torch.float64)),
 'DR006': ('Drift',
           tensor(0.1096, dtype=torch.float64),
           tensor(0., dtype=torch.float64)),
 'DR007': ('Drift',
           tensor(0.0440, dtype=torch.float64),
           tensor(0., dtype=torch.float64)),
 'RM5': ('Dipole',
         tensor(0.8728, dtype=torch.float64),
         tensor(0.7854, dtype=torch.float64)),
 'SY1_3F4': ('Sextupole',
             tensor(0.0800, dtype=torch.float64),
             tensor(0., dtype=torch.float64))}
[44]:
# duplicate (check first level for duplicates (elements with the same name)

line, *_ = ring.clone()
print(line.duplicates)

_, line, *_ = ring.clone()
print(line.duplicates)
True
False
[45]:
# itemize (get element names with matching kind at all levels, ordered, with duplicates)

pprint(ring.itemize('Dipole'))
pprint(ring.itemize('Quadrupole'))
['RM5', 'RM6', 'RM7', 'RM8', 'RM1', 'RM2', 'RM3', 'RM4']
['Q3D3',
 'Q3D2',
 'Q3F2',
 'Q3F4',
 'Q3F1',
 'Q3D1',
 'Q4D1',
 'Q4F1',
 'Q4F4',
 'Q4F2',
 'Q4D2',
 'Q4D3',
 'Q4F3',
 'Q1F3',
 'Q1D3',
 'Q1D2',
 'Q1F2',
 'Q1F4',
 'Q1F1',
 'Q1D1',
 'Q2D1',
 'Q2F1',
 'Q2F4',
 'Q2F2',
 'Q2D2',
 'Q2D3',
 'Q2F3',
 'Q3F3']
[46]:
# describe (element counts by kinds, all levels, with duplicates)

pprint(ring.describe)
{'BPM': 16,
 'Corrector': 28,
 'Dipole': 8,
 'Drift': 67,
 'Quadrupole': 28,
 'Sextupole': 16}
[47]:
# split (split elements)

_, line, *_ = ring.clone()
print(line)
print()

# zero or one do not slice matching element(s)

_, line, *_ = ring.clone()
line.split((1, ['Dipole'], None, None))
print(line)
print()

# split by kind

_, line, *_ = ring.clone()

line.split((4, ['Dipole'], None, None))
print(line)
print()

# split by kind with insertion (number of insertions is count - 1)

_, line, *_ = ring.clone()

line.split((4, ['Dipole'], None, None), paste=[Corrector('CXY')])
print(line)
print()


# split by names

_, line, *_ = ring.clone()

line.split((4, None, ['RM5'], None))
print(line)
print()

# split by names with exclude

_, line, *_ = ring.clone()

line.split((4, None, ['RM5'], ['RM5']))
print(line)
print()
BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")

BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")

BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.21821, angle=0.19634954085000025, e1=0.0, e1_on=True, e2=0.0, e2_on=False, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.21821, angle=0.19634954085000025, e1=0.0, e1_on=False, e2=0.0, e2_on=False, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.21821, angle=0.19634954085000025, e1=0.0, e1_on=False, e2=0.0, e2_on=False, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.21821, angle=0.19634954085000025, e1=0.0, e1_on=False, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")

BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.21821, angle=0.19634954085000025, e1=0.0, e1_on=True, e2=0.0, e2_on=False, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Dipole(name="RM5", length=0.21821, angle=0.19634954085000025, e1=0.0, e1_on=False, e2=0.0, e2_on=False, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Dipole(name="RM5", length=0.21821, angle=0.19634954085000025, e1=0.0, e1_on=False, e2=0.0, e2_on=False, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY", cx=0.0, cy=0.0)
Dipole(name="RM5", length=0.21821, angle=0.19634954085000025, e1=0.0, e1_on=False, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")

BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.21821, angle=0.19634954085000025, e1=0.0, e1_on=True, e2=0.0, e2_on=False, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.21821, angle=0.19634954085000025, e1=0.0, e1_on=False, e2=0.0, e2_on=False, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.21821, angle=0.19634954085000025, e1=0.0, e1_on=False, e2=0.0, e2_on=False, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.21821, angle=0.19634954085000025, e1=0.0, e1_on=False, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")

BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")

[48]:
# clean (remove elements by kind, name, length)
# Note, it element has zero length, it will not be removed

_, line, *_ = ring.clone()
print(line)
print()

# Remove BPMs

line.clean((None, ['BPM'], None, None))
print(line)
print()

# Remove by name

line.clean((None, None, ['RM5'], None))
print(line)
print()

# Remove by length

line.clean((0.11, None, None, None))
print(line)
print()
BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")

Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)

Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)

Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)

[49]:
# merge (merge drifts)
# Note, all drift elements are renamed

_, line, *_ = ring.clone()
line.remove("RM5")
print(line)
print()

line.merge()
print(line)
print()
BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")

BPM(name="BPM07", direction="inverse")
Drift(name="DR001", length=0.22916, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR002", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")

[50]:
# mangle (rename elements with identical names)

line, *_ = ring.clone()
print(line)
print()

# Rename quadrupoles

line, *_ = ring.clone()
line.mangle('Quadrupole')
print(line)
print()

# Rename quadrupoles with exclude

line, *_ = ring.clone()
line.mangle('Quadrupole', names=['Q3D3'])
print(line)
print()
BPM(name="BPM05", direction="inverse")
Drift(name="DR001", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D3", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_001", cx=0.0, cy=0.0)
Quadrupole(name="Q3D3", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR002", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D2", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_002", cx=0.0, cy=0.0)
Quadrupole(name="Q3D2", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR003", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F2", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_003", cx=0.0, cy=0.0)
Quadrupole(name="Q3F2", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR004", length=0.065, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM07", direction="forward")

BPM(name="BPM05", direction="inverse")
Drift(name="DR001", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D3_001", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_001", cx=0.0, cy=0.0)
Quadrupole(name="Q3D3_002", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR002", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D2_001", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_002", cx=0.0, cy=0.0)
Quadrupole(name="Q3D2_002", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR003", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F2_001", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_003", cx=0.0, cy=0.0)
Quadrupole(name="Q3F2_002", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR004", length=0.065, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM07", direction="forward")

BPM(name="BPM05", direction="inverse")
Drift(name="DR001", length=0.116715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D3", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_001", cx=0.0, cy=0.0)
Quadrupole(name="Q3D3", length=0.098285, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR002", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3D2_001", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_002", cx=0.0, cy=0.0)
Quadrupole(name="Q3D2_002", length=0.098285, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR003", length=0.221715, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="Q3F2_001", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Corrector(name="CXY_003", cx=0.0, cy=0.0)
Quadrupole(name="Q3F2_002", length=0.09, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR004", length=0.065, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM07", direction="forward")

[51]:
# splice (sequencies between BPMs)

line = ring.clone()

line.flatten()
print(len(line))

line.splice()
print(len(line))
207
16
[52]:
# dp (get/set)

_, line, *_ = ring.clone()
line.propagate = True
line.dp = 0.001
print(line.dp)
print()

print(line)
tensor(0.0010, dtype=torch.float64)

BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.001, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.001, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.001, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.001, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.001, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")
[53]:
# exact (get/set)

_, line, *_ = ring.clone()
line.propagate = True
line.exact = True
print(line)
BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.0, exact=True, ns=1, order=0)
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=True, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=True, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=True, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=True, ns=1, order=0)
BPM(name="BPM08", direction="forward")
[54]:
# output (get/set)

_, line, *_ = ring.clone()
line.propagate = True
line.output = True
line['RM5'].output
[54]:
True
[55]:
# matrix (get/set)

_, line, *_ = ring.clone()
line.propagate = True
line.matrix = True
line['RM5'].matrix
[55]:
True
[56]:
# length

line = ring.clone()
print(line.length)
tensor(27.3689, dtype=torch.float64)
[57]:
# angle

line = ring.clone()
print(line.angle)
tensor(6.2832, dtype=torch.float64)
[58]:
# flag (layout flag, true if contains dipoles with non-zero angles)

line = ring.clone()
print(line.flag)
True
[59]:
# ns (get/set number of integration steps)

_, line, *_ = ring.clone()
line.propagate = True
pprint(line.ns)
print()

# Set fixed number to all elements

_, line, *_ = ring.clone()
line.propagate = True
line.ns = 10
pprint(line.ns)
print()

# Set by step length (keeps 1 step for zero length elements)

_, line, *_ = ring.clone()
line.propagate = True
line.ns = 0.01
pprint(line.ns)
print()

# Set by kind (pass int or float)

_, line, *_ = ring.clone()
line.propagate = True
line.ns = (('Dipole', 0.01), ('Sextupole', 0.01))
pprint(line.ns)
print()

# Set by name (pass int or float)

_, line, *_ = ring.clone()
line.propagate = True
line.ns = (('RM5', 10),)
pprint(line.ns)
print()
{'BPM07': 1,
 'BPM08': 1,
 'DR005': 1,
 'DR006': 1,
 'DR007': 1,
 'RM5': 1,
 'SY1_3F4': 1}

{'BPM07': 10,
 'BPM08': 10,
 'DR005': 10,
 'DR006': 10,
 'DR007': 10,
 'RM5': 10,
 'SY1_3F4': 10}

{'BPM07': 1,
 'BPM08': 1,
 'DR005': 12,
 'DR006': 11,
 'DR007': 5,
 'RM5': 88,
 'SY1_3F4': 8}

{'BPM07': 1,
 'BPM08': 1,
 'DR005': 1,
 'DR006': 1,
 'DR007': 1,
 'RM5': 88,
 'SY1_3F4': 8}

{'BPM07': 1,
 'BPM08': 1,
 'DR005': 1,
 'DR006': 1,
 'DR007': 1,
 'RM5': 10,
 'SY1_3F4': 1}

[60]:
# order (get/set integration order)

_, line, *_ = ring.clone()
line.propagate = True
pprint(line.order)
print()

# Set fixed number to all elements

_, line, *_ = ring.clone()
line.propagate = True
line.order = 1
pprint(line.order)
print()

# Set by kind

_, line, *_ = ring.clone()
line.propagate = True
line.order = (('Dipole', 1), ('Sextupole', 1))
pprint(line.order)
print()

# Set by name

_, line, *_ = ring.clone()
line.propagate = True
line.order = (('RM5', 1),)
pprint(line.order)
print()
{'BPM07': 0,
 'BPM08': 0,
 'DR005': 0,
 'DR006': 0,
 'DR007': 0,
 'RM5': 0,
 'SY1_3F4': 0}

{'BPM07': 1,
 'BPM08': 1,
 'DR005': 1,
 'DR006': 1,
 'DR007': 1,
 'RM5': 1,
 'SY1_3F4': 1}

{'BPM07': 0,
 'BPM08': 0,
 'DR005': 0,
 'DR006': 0,
 'DR007': 0,
 'RM5': 1,
 'SY1_3F4': 1}

{'BPM07': 0,
 'BPM08': 0,
 'DR005': 0,
 'DR006': 0,
 'DR007': 0,
 'RM5': 1,
 'SY1_3F4': 0}

[61]:
# __call__ (propagate initial condition, pass deviation variables)

line:Line = ring.clone()
line.propagate = True

# Propagate initial condition

state = torch.tensor([0.001, 0.0005, -0.005, 0.0001], dtype=torch.float64)
print(line(state))

# Equivalent to element by element propagation of a flat line

line.flatten()
state = torch.tensor([0.001, 0.0005, -0.005, 0.0001], dtype=torch.float64)
for element in line.sequence:
    state = element(state)
print(state)

# Passing deviation parameters

line:Line = ring.clone()
line.propagate = True
line.flatten()

data = line.data()
data['Q3D3']['kn'] = torch.tensor(0.1, dtype=torch.float64)
state = torch.tensor([0.001, 0.0005, -0.005, 0.0001], dtype=torch.float64)
print(line(state, data=data))

# Equivalent to changing kn value in all Q3D3 occurrancies

for position in line.positions('Q3D3'):
    line[position].kn = line[position].kn.item() + 0.1

state = torch.tensor([0.001, 0.0005, -0.005, 0.0001], dtype=torch.float64)
print(line(state))

# Passing alignment errors as deviation parameters

line:Line = ring.clone()
line.propagate = True
line.flatten()

data = line.data()
data['Q3D3']['dx'] = torch.tensor(0.001, dtype=torch.float64)
state = torch.tensor([0.001, 0.0005, -0.005, 0.0001], dtype=torch.float64)
print(line(state, data=data, alignment=True))

# Equivalent to changing dx value in all Q3D3 occurrancies

for position in line.positions('Q3D3'):
    line[position].dx = 0.001

state = torch.tensor([0.001, 0.0005, -0.005, 0.0001], dtype=torch.float64)
print(line(state, alignment=True))
tensor([ 0.0016, -0.0031, -0.0193, -0.0315], dtype=torch.float64)
tensor([ 0.0016, -0.0031, -0.0193, -0.0315], dtype=torch.float64)
tensor([ 0.0015, -0.0028, -0.0190, -0.0310], dtype=torch.float64)
tensor([ 0.0015, -0.0028, -0.0190, -0.0310], dtype=torch.float64)
tensor([ 0.0048, -0.0084, -0.0187, -0.0302], dtype=torch.float64)
tensor([ 0.0048, -0.0084, -0.0187, -0.0302], dtype=torch.float64)
[62]:
# __len__ (number of first level elements in the sequence)

line:Line = ring.clone()
line.propagate = True
line.flatten()

print(len(line))
print(len(line.sequence))
207
207
[63]:
# layout  (name, kind, length, angle) for all leaf elements (ordered)

_, line, *_ = ring.clone()

pprint(line.layout())
[('BPM07',
  'BPM',
  tensor(0., dtype=torch.float64),
  tensor(0., dtype=torch.float64)),
 ('DR005',
  'Drift',
  tensor(0.1196, dtype=torch.float64),
  tensor(0., dtype=torch.float64)),
 ('RM5',
  'Dipole',
  tensor(0.8728, dtype=torch.float64),
  tensor(0.7854, dtype=torch.float64)),
 ('DR006',
  'Drift',
  tensor(0.1096, dtype=torch.float64),
  tensor(0., dtype=torch.float64)),
 ('SY1_3F4',
  'Sextupole',
  tensor(0.0800, dtype=torch.float64),
  tensor(0., dtype=torch.float64)),
 ('DR007',
  'Drift',
  tensor(0.0440, dtype=torch.float64),
  tensor(0., dtype=torch.float64)),
 ('BPM08',
  'BPM',
  tensor(0., dtype=torch.float64),
  tensor(0., dtype=torch.float64))]
[64]:
# __repr__ (string representation, can be used for instance creation, not recommended)

_, line, *_ = ring.clone()

print(line)
print()

element, *_ = line

print(element)
print()

from model.library.bpm import BPM
print(eval(str(element)))
BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")

BPM(name="BPM07", direction="inverse")

BPM(name="BPM07", direction="inverse")
[65]:
# __getitem__ and __setitem__ (similar for __delitem__)

# Supports indexing by name (first match) or index in the sequence
# Supports indexing of nested lines
# Unpacking also works

line = ring.clone()

print(line[1])
print()

print(line[1, 0])
print()

print(line['BPM07_BPM08', 'BPM07'])
print()

line['BPM07_BPM08', 'BPM07'] = Corrector('CXY')
print(line[1])
print()
BPM(name="BPM07", direction="inverse")
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")

BPM(name="BPM07", direction="inverse")

BPM(name="BPM07", direction="inverse")

Corrector(name="CXY", cx=0.0, cy=0.0)
Drift(name="DR005", length=0.11958, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="RM5", length=0.87284, angle=0.785398163400001, e1=0.0, e1_on=True, e2=0.0, e2_on=True, kn=-2.379107171999999, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR006", length=0.10958, dp=0.0, exact=False, ns=1, order=0)
Sextupole(name="SY1_3F4", length=0.08, ms=-277.23165, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR007", length=0.044, dp=0.0, exact=False, ns=1, order=0)
BPM(name="BPM08", direction="forward")

Example-22: Wrapper

[1]:
# In this example construction of parametric call wrappers is illustrated
# For elements, all deviation parameters are passed as dictionary
# Wrapped elements are invoked using positional agruments
[2]:
# Import

from pprint import pprint

import torch

from twiss import twiss

from ndmap.pfp import parametric_fixed_point
from ndmap.evaluate import evaluate
from ndmap.signature import chop

from model.library.drift import Drift
from model.library.multipole import Multipole
from model.library.dipole import Dipole
from model.library.line import Line

from model.command.wrapper import wrapper
[3]:
# Define simple FODO based lattice using nested lines

QF = Multipole('QF', 0.5, +0.20)
QD = Multipole('QD', 0.5, -0.19)
DR = Drift('DR', 0.75)
BM = Dipole('BM', 3.50, torch.pi/4.0)

FODO = [QF, DR, BM, DR, QD, QD, DR, BM, DR, QF]

FODO_A = Line('FODO_A', FODO, propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_B = Line('FODO_B', FODO, propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_C = Line('FODO_C', FODO, propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_D = Line('FODO_D', FODO, propagate=True, dp=0.0, exact=False, output=False, matrix=False)

LINE_AB = Line('LINE_AB', [FODO_A, FODO_B], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
LINE_CD = Line('LINE_CD', [FODO_C, FODO_D], propagate=True, dp=0.0, exact=False, output=False, matrix=False)

RING = Line('RING', [LINE_AB, LINE_CD], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
[4]:
# Deviation variables are passed to elements/lines as dictionaries
# In order to compute derivatives with respect to a deviation variable
# A tensor should be binded to a corresponding leaf deviation dictionary value

pprint(RING.data(alignment=False), sort_dicts=False)
{'LINE_AB': {'FODO_A': {'QF': {'kn': tensor(0., dtype=torch.float64),
                               'ks': tensor(0., dtype=torch.float64),
                               'ms': tensor(0., dtype=torch.float64),
                               'mo': tensor(0., dtype=torch.float64),
                               'dp': tensor(0., dtype=torch.float64),
                               'dl': tensor(0., dtype=torch.float64)},
                        'DR': {'dp': tensor(0., dtype=torch.float64),
                               'dl': tensor(0., dtype=torch.float64)},
                        'BM': {'dw': tensor(0., dtype=torch.float64),
                               'e1': tensor(0., dtype=torch.float64),
                               'e2': tensor(0., dtype=torch.float64),
                               'kn': tensor(0., dtype=torch.float64),
                               'ks': tensor(0., dtype=torch.float64),
                               'ms': tensor(0., dtype=torch.float64),
                               'mo': tensor(0., dtype=torch.float64),
                               'dp': tensor(0., dtype=torch.float64),
                               'dl': tensor(0., dtype=torch.float64)},
                        'QD': {'kn': tensor(0., dtype=torch.float64),
                               'ks': tensor(0., dtype=torch.float64),
                               'ms': tensor(0., dtype=torch.float64),
                               'mo': tensor(0., dtype=torch.float64),
                               'dp': tensor(0., dtype=torch.float64),
                               'dl': tensor(0., dtype=torch.float64)}},
             'FODO_B': {'QF': {'kn': tensor(0., dtype=torch.float64),
                               'ks': tensor(0., dtype=torch.float64),
                               'ms': tensor(0., dtype=torch.float64),
                               'mo': tensor(0., dtype=torch.float64),
                               'dp': tensor(0., dtype=torch.float64),
                               'dl': tensor(0., dtype=torch.float64)},
                        'DR': {'dp': tensor(0., dtype=torch.float64),
                               'dl': tensor(0., dtype=torch.float64)},
                        'BM': {'dw': tensor(0., dtype=torch.float64),
                               'e1': tensor(0., dtype=torch.float64),
                               'e2': tensor(0., dtype=torch.float64),
                               'kn': tensor(0., dtype=torch.float64),
                               'ks': tensor(0., dtype=torch.float64),
                               'ms': tensor(0., dtype=torch.float64),
                               'mo': tensor(0., dtype=torch.float64),
                               'dp': tensor(0., dtype=torch.float64),
                               'dl': tensor(0., dtype=torch.float64)},
                        'QD': {'kn': tensor(0., dtype=torch.float64),
                               'ks': tensor(0., dtype=torch.float64),
                               'ms': tensor(0., dtype=torch.float64),
                               'mo': tensor(0., dtype=torch.float64),
                               'dp': tensor(0., dtype=torch.float64),
                               'dl': tensor(0., dtype=torch.float64)}}},
 'LINE_CD': {'FODO_C': {'QF': {'kn': tensor(0., dtype=torch.float64),
                               'ks': tensor(0., dtype=torch.float64),
                               'ms': tensor(0., dtype=torch.float64),
                               'mo': tensor(0., dtype=torch.float64),
                               'dp': tensor(0., dtype=torch.float64),
                               'dl': tensor(0., dtype=torch.float64)},
                        'DR': {'dp': tensor(0., dtype=torch.float64),
                               'dl': tensor(0., dtype=torch.float64)},
                        'BM': {'dw': tensor(0., dtype=torch.float64),
                               'e1': tensor(0., dtype=torch.float64),
                               'e2': tensor(0., dtype=torch.float64),
                               'kn': tensor(0., dtype=torch.float64),
                               'ks': tensor(0., dtype=torch.float64),
                               'ms': tensor(0., dtype=torch.float64),
                               'mo': tensor(0., dtype=torch.float64),
                               'dp': tensor(0., dtype=torch.float64),
                               'dl': tensor(0., dtype=torch.float64)},
                        'QD': {'kn': tensor(0., dtype=torch.float64),
                               'ks': tensor(0., dtype=torch.float64),
                               'ms': tensor(0., dtype=torch.float64),
                               'mo': tensor(0., dtype=torch.float64),
                               'dp': tensor(0., dtype=torch.float64),
                               'dl': tensor(0., dtype=torch.float64)}},
             'FODO_D': {'QF': {'kn': tensor(0., dtype=torch.float64),
                               'ks': tensor(0., dtype=torch.float64),
                               'ms': tensor(0., dtype=torch.float64),
                               'mo': tensor(0., dtype=torch.float64),
                               'dp': tensor(0., dtype=torch.float64),
                               'dl': tensor(0., dtype=torch.float64)},
                        'DR': {'dp': tensor(0., dtype=torch.float64),
                               'dl': tensor(0., dtype=torch.float64)},
                        'BM': {'dw': tensor(0., dtype=torch.float64),
                               'e1': tensor(0., dtype=torch.float64),
                               'e2': tensor(0., dtype=torch.float64),
                               'kn': tensor(0., dtype=torch.float64),
                               'ks': tensor(0., dtype=torch.float64),
                               'ms': tensor(0., dtype=torch.float64),
                               'mo': tensor(0., dtype=torch.float64),
                               'dp': tensor(0., dtype=torch.float64),
                               'dl': tensor(0., dtype=torch.float64)},
                        'QD': {'kn': tensor(0., dtype=torch.float64),
                               'ks': tensor(0., dtype=torch.float64),
                               'ms': tensor(0., dtype=torch.float64),
                               'mo': tensor(0., dtype=torch.float64),
                               'dp': tensor(0., dtype=torch.float64),
                               'dl': tensor(0., dtype=torch.float64)}}}}
[5]:
# Compute parametric closed orbit (first order with respect to momentum deviation)

# Without wrapping, all momenta deviation occurances should be binded to a singel tensor
# Hence, deviation table should be traversed recursively down to all leafs

def scan(data, name, target):
    for key, value in data.items():
        if isinstance(value, dict):
            scan(value, name, target)
        elif key == name:
            data[key] = target

# Set ring function

def ring(state, dp):
    dp, *_ = dp
    data = RING.data()
    scan(data, 'dp', dp)
    return  RING(state, data=data)

# Set deviations

fp = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)
dp = torch.tensor([0.0], dtype=torch.float64)

# Compute pfp

pfp, *_ = parametric_fixed_point((1, ), fp, [dp], ring)
chop(pfp)
pfp
[5]:
[tensor([0., 0., 0., 0.], dtype=torch.float64),
 tensor([[4.4462],
         [0.0000],
         [0.0000],
         [0.0000]], dtype=torch.float64)]
[6]:
# Using wrapper we can define the about ring function as follows

fn = wrapper(RING, (None, None, 'dp'))

# Set deviations

fp = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)
dp = torch.tensor([0.0], dtype=torch.float64)

# Compute pfp

pfp, *_ = parametric_fixed_point((1, ), fp, [dp], fn)
chop(pfp)
pfp
[6]:
[tensor([0., 0., 0., 0.], dtype=torch.float64),
 tensor([[4.4462],
         [0.0000],
         [0.0000],
         [0.0000]], dtype=torch.float64)]
[7]:
# Compute chromaticity (without wrapping)

def scan(data, name, target):
    for key, value in data.items():
        if isinstance(value, dict):
            scan(value, name, target)
        elif key == name:
            data[key] = target

# Set ring function

def ring(state, dp):
    dp, *_ = dp
    data = RING.data()
    scan(data, 'dp', dp)
    return RING(state , data=data)

# Set ring function around pfp

def pfp_ring(state, dp):
    return ring(state + evaluate(pfp, [dp]), dp) - evaluate(pfp, [dp])

# Set tune function

def tune(dp):
    matrix = torch.func.jacrev(pfp_ring)(state, dp)
    tunes, *_ = twiss(matrix)
    return tunes

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)
dp = torch.tensor([0.0], dtype=torch.float64)

print(tune(dp))
print(torch.func.jacrev(tune)(dp).squeeze())
tensor([0.6951, 0.7019], dtype=torch.float64)
tensor([-2.0649, -0.8260], dtype=torch.float64)
[8]:
# Compute chromaticity (with wrapping)

# Set ring function

fn = wrapper(RING, (None, None, 'dp'))

# Set ring function around pfp

def pfp_ring(state, dp):
    return fn(state + evaluate(pfp, [dp]), dp) - evaluate(pfp, [dp])

# Set tune function

def tune(dp):
    matrix = torch.func.jacrev(pfp_ring)(state, dp)
    tunes, *_ = twiss(matrix)
    return tunes

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)
dp = torch.tensor([0.0], dtype=torch.float64)

print(tune(dp))
print(torch.func.jacrev(tune)(dp).squeeze())
tensor([0.6951, 0.7019], dtype=torch.float64)
tensor([-2.0649, -0.8260], dtype=torch.float64)
[9]:
# Compute chromaticity derivative with respect to sextupole ampitudes (without wrapping)

def scan(data, name, target):
    for key, value in data.items():
        if isinstance(value, dict):
            scan(value, name, target)
        elif key == name:
            data[key] = target

def ring(state, dp, dms):
    dp, *_ = dp
    dmsf, dmsd, *_ = dms
    data = RING.data()
    scan(data, 'dp', dp)
    data['LINE_AB']['FODO_A']['QF']['ms'] = dmsf
    data['LINE_AB']['FODO_B']['QF']['ms'] = dmsf
    data['LINE_CD']['FODO_C']['QF']['ms'] = dmsf
    data['LINE_CD']['FODO_D']['QF']['ms'] = dmsf
    data['LINE_AB']['FODO_A']['QD']['ms'] = dmsd
    data['LINE_AB']['FODO_B']['QD']['ms'] = dmsd
    data['LINE_CD']['FODO_C']['QD']['ms'] = dmsd
    data['LINE_CD']['FODO_D']['QD']['ms'] = dmsd
    return RING(state, data=data)

def pfp_ring(state, dp, dms):
    return ring(state + evaluate(pfp, [dp]), dp, dms) - evaluate(pfp, [dp])


def tune(dp, dms):
    matrix = torch.func.jacrev(pfp_ring)(state, dp, dms)
    tunes, *_ = twiss(matrix)
    return tunes

def chromaticity(dms):
    return torch.func.jacrev(tune)(dp, dms)


state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)
dms = torch.tensor([0.0, 0.0], dtype=torch.float64)
dp = torch.tensor([0.0], dtype=torch.float64)

print(tune(dp, dms))
print(chromaticity(dms).squeeze())
print(torch.func.jacrev(chromaticity)(dms).squeeze())
tensor([0.6951, 0.7019], dtype=torch.float64)
tensor([-2.0649, -0.8260], dtype=torch.float64)
tensor([[ 25.8500,   1.0470],
        [ -9.0271, -16.4821]], dtype=torch.float64)
[10]:
# Compute chromaticity derivative with respect to sextupole ampitudes (with wrapping)

def scan(data, name, target):
    for key, value in data.items():
        if isinstance(value, dict):
            scan(value, name, target)
        elif key == name:
            data[key] = target

ring = wrapper(RING, (None, None, 'dp'), (None, ['QF', 'QD'], 'ms'))

def pfp_ring(state, dp, dms):
    return ring(state + evaluate(pfp, [dp]), dp, dms) - evaluate(pfp, [dp])


def tune(dp, dms):
    matrix = torch.func.jacrev(pfp_ring)(state, dp, dms)
    tunes, *_ = twiss(matrix)
    return tunes

def chromaticity(dms):
    return torch.func.jacrev(tune)(dp, dms)


state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)
dms = torch.tensor([0.0, 0.0], dtype=torch.float64)
dp = torch.tensor([0.0], dtype=torch.float64)

print(tune(dp, dms))
print(chromaticity(dms).squeeze())
print(torch.func.jacrev(chromaticity)(dms).squeeze())
tensor([0.6951, 0.7019], dtype=torch.float64)
tensor([-2.0649, -0.8260], dtype=torch.float64)
tensor([[ 25.8500,   1.0470],
        [ -9.0271, -16.4821]], dtype=torch.float64)
[11]:
# The above examples demonstrate how to bind tensors to all leafs or to given elements (in all lines)

# (None,           None,            parameter:str) -- bind tensor to all leaf parameters
# (None,           names:list[str], parameter:str) -- bind tensor to all leaf parameters in specified elements
# (path:list[str], names:list[str], parameter:str) -- bind tensor to all leaf parameters in specified elements in given path (path to specific line)
[12]:
# Bind QF and QD in all sublines of a given line (1/2 of sextupoles)

ring = wrapper(RING, (None, None, 'dp'), (['LINE_AB'], ['QF', 'QD'], 'ms'))

def pfp_ring(state, dp, dms):
    return ring(state + evaluate(pfp, [dp]), dp, dms) - evaluate(pfp, [dp])

def tune(dp, dms):
    matrix = torch.func.jacrev(pfp_ring)(state, dp, dms)
    tunes, *_ = twiss(matrix)
    return tunes

def chromaticity(dms):
    return torch.func.jacrev(tune)(dp, dms)

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)
dms = torch.tensor([0.0, 0.0], dtype=torch.float64)
dp = torch.tensor([0.0], dtype=torch.float64)

print(tune(dp, dms))
print(chromaticity(dms).squeeze())
print(2*torch.func.jacrev(chromaticity)(dms).squeeze())
tensor([0.6951, 0.7019], dtype=torch.float64)
tensor([-2.0649, -0.8260], dtype=torch.float64)
tensor([[ 25.8500,   1.0470],
        [ -9.0271, -16.4821]], dtype=torch.float64)
[13]:
# Bind QF and QD in a given leaf line (1/4 of sextupoles)

ring = wrapper(RING, (None, None, 'dp'), (['LINE_AB', 'FODO_A'], ['QF', 'QD'], 'ms'))

def pfp_ring(state, dp, dms):
    return ring(state + evaluate(pfp, [dp]), dp, dms) - evaluate(pfp, [dp])

def tune(dp, dms):
    matrix = torch.func.jacrev(pfp_ring)(state, dp, dms)
    tunes, *_ = twiss(matrix)
    return tunes

def chromaticity(dms):
    return torch.func.jacrev(tune)(dp, dms)

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)
dms = torch.tensor([0.0, 0.0], dtype=torch.float64)
dp = torch.tensor([0.0], dtype=torch.float64)

print(tune(dp, dms))
print(chromaticity(dms).squeeze())
print(4*torch.func.jacrev(chromaticity)(dms).squeeze())
tensor([0.6951, 0.7019], dtype=torch.float64)
tensor([-2.0649, -0.8260], dtype=torch.float64)
tensor([[ 25.8500,   1.0470],
        [ -9.0271, -16.4821]], dtype=torch.float64)
[14]:
# Several sextupole groups

ring = wrapper(RING, (None, None, 'dp'), (['LINE_AB'], ['QF', 'QD'], 'ms'), (['LINE_CD'], ['QF', 'QD'], 'ms'))

def pfp_ring(state, dp, dms_ab, dms_cd):
    return ring(state + evaluate(pfp, [dp]), dp, dms_ab, dms_cd) - evaluate(pfp, [dp])

def tune(dp, dms_ab, dms_cd):
    matrix = torch.func.jacrev(pfp_ring)(state, dp, dms_ab, dms_cd)
    tunes, *_ = twiss(matrix)
    return tunes

def chromaticity(dms_ab, dms_cd):
    return torch.func.jacrev(tune)(dp, dms_ab, dms_cd)

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)
dms_ab = torch.tensor([0.0, 0.0], dtype=torch.float64)
dms_cd = torch.tensor([0.0, 0.0], dtype=torch.float64)
dp = torch.tensor([0.0], dtype=torch.float64)

print(tune(dp, dms_ab, dms_cd))
print(chromaticity(dms_ab, dms_cd).squeeze())
print(torch.func.jacrev(chromaticity, 0)(dms_ab, dms_cd).squeeze())
print(torch.func.jacrev(chromaticity, 1)(dms_ab, dms_cd).squeeze())

def fn(dms):
    dms_ab, dms_cd = dms
    return chromaticity(dms_ab, dms_cd)

print(torch.func.jacrev(fn)(torch.stack([dms_ab, dms_cd])).squeeze())
tensor([0.6951, 0.7019], dtype=torch.float64)
tensor([-2.0649, -0.8260], dtype=torch.float64)
tensor([[12.9250,  0.5235],
        [-4.5135, -8.2411]], dtype=torch.float64)
tensor([[12.9250,  0.5235],
        [-4.5135, -8.2411]], dtype=torch.float64)
tensor([[[12.9250,  0.5235],
         [12.9250,  0.5235]],

        [[-4.5135, -8.2411],
         [-4.5135, -8.2411]]], dtype=torch.float64)

Example-23: Group

[1]:
# In this example another wrapper construction procedure is illustraded
[2]:
# Import

import torch
torch.set_printoptions(linewidth=128)

from twiss import twiss
from twiss import propagate
from twiss import wolski_to_cs

from model.library.drift import Drift
from model.library.quadrupole import Quadrupole
from model.library.dipole import Dipole
from model.library.line import Line

from model.command.wrapper import group
[3]:
# Define simple FODO based lattice using nested lines

DR = Drift('DR', 0.75)
BM = Dipole('BM', 3.50, torch.pi/4.0)

QF_A = Quadrupole('QF_A', 0.5, +0.20)
QD_A = Quadrupole('QD_A', 0.5, -0.19)
QF_B = Quadrupole('QF_B', 0.5, +0.20)
QD_B = Quadrupole('QD_B', 0.5, -0.19)
QF_C = Quadrupole('QF_C', 0.5, +0.20)
QD_C = Quadrupole('QD_C', 0.5, -0.19)
QF_D = Quadrupole('QF_D', 0.5, +0.20)
QD_D = Quadrupole('QD_D', 0.5, -0.19)

FODO_A = Line('FODO_A', [QF_A, DR, BM, DR, QD_A, QD_A, DR, BM, DR, QF_A], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_B = Line('FODO_B', [QF_B, DR, BM, DR, QD_B, QD_B, DR, BM, DR, QF_B], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_C = Line('FODO_C', [QF_C, DR, BM, DR, QD_C, QD_C, DR, BM, DR, QF_C], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_D = Line('FODO_D', [QF_D, DR, BM, DR, QD_D, QD_D, DR, BM, DR, QF_D], propagate=True, dp=0.0, exact=False, output=False, matrix=False)

RING = Line('RING', [FODO_A, FODO_B, FODO_C, FODO_D], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
[4]:
# Full ring or subline can be wrapped by element kind

fn, table, line = group(RING,                               # -- source line
                        'FODO_A',                           # -- start (name or position in source line sequence)
                        'FODO_B',                           # -- end (name or position in source line sequence)
                        ('kn', ['Quadrupole'], None, None)) # -- groups (key:str, kinds:list[str]|None, names:list[str]|None, clean:list[str]|None

# Information about deviation variables is returbed in wrapper format

print(table)
print()

# Wrapped function fn can be called with deviation variables

(_, names, _), *_ = table
knobs = torch.tensor(len(names)*[0.0], dtype=torch.float64)
state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)
print(fn(state, knobs))
print()

# Constructed line also returned

print(line)
print()
[(None, ['QF_A', 'QD_A', 'QF_B', 'QD_B'], 'kn')]

tensor([0., 0., 0., 0.], dtype=torch.float64)

Quadrupole(name="QF_A", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR", length=0.75, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="BM", length=3.5, angle=0.7853981633974493, e1=0.0, e2=0.0, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR", length=0.75, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="QD_A", length=0.5, kn=-0.189999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="QD_A", length=0.5, kn=-0.189999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR", length=0.75, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="BM", length=3.5, angle=0.7853981633974493, e1=0.0, e2=0.0, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR", length=0.75, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="QF_A", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="QF_B", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR", length=0.75, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="BM", length=3.5, angle=0.7853981633974493, e1=0.0, e2=0.0, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR", length=0.75, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="QD_B", length=0.5, kn=-0.189999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="QD_B", length=0.5, kn=-0.189999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR", length=0.75, dp=0.0, exact=False, ns=1, order=0)
Dipole(name="BM", length=3.5, angle=0.7853981633974493, e1=0.0, e2=0.0, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)
Drift(name="DR", length=0.75, dp=0.0, exact=False, ns=1, order=0)
Quadrupole(name="QF_B", length=0.5, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)

[5]:
# By default names are excracted from created subline

_, table, _ = group(RING, 'FODO_A', 'FODO_B', ('kn', ['Quadrupole'], None, None))

print(table)
print()

# Use root flag to extract name from the root line


_, table, _ = group(RING, 'FODO_A', 'FODO_B', ('kn', ['Quadrupole'], None, None), root=True)

print(table)
print()
[(None, ['QF_A', 'QD_A', 'QF_B', 'QD_B'], 'kn')]

[(None, ['QF_A', 'QD_A', 'QF_B', 'QD_B', 'QF_C', 'QD_C', 'QF_D', 'QD_D'], 'kn')]

[6]:
# Set transport between observation points

# 0--A--1--B--2--C--3--D--4

line01, *_ =  group(RING, 'FODO_A', 'FODO_A', ('kn', ['Quadrupole'], None, None), root=True)
line12, *_ =  group(RING, 'FODO_B', 'FODO_B', ('kn', ['Quadrupole'], None, None), root=True)
line23, *_ =  group(RING, 'FODO_C', 'FODO_C', ('kn', ['Quadrupole'], None, None), root=True)
line34, *_ =  group(RING, 'FODO_D', 'FODO_D', ('kn', ['Quadrupole'], None, None), root=True)

lines = [
    line01,
    line12,
    line23,
    line34
]

def ring(state, knobs):
    for line in lines:
        state = line(state, knobs)
    return state


state = torch.tensor(4*[0.0], dtype=torch.float64)
knobs = torch.tensor(8*[0.0], dtype=torch.float64)

print(ring(state, knobs))
tensor([0., 0., 0., 0.], dtype=torch.float64)
[7]:
# Compute tunes and corresponding derivatives with respect to deviation parameters

def fn(knobs):
    m = torch.func.jacfwd(ring)(state, knobs)
    t, *_ = twiss(m)
    return t

print(fn(knobs))
print(torch.func.jacrev(fn)(knobs))
tensor([0.6951, 0.7019], dtype=torch.float64)
tensor([[ 1.4567,  0.1055,  1.4567,  0.1055,  1.4567,  0.1055,  1.4567,  0.1055],
        [-0.5132, -1.6271, -0.5132, -1.6271, -0.5132, -1.6271, -0.5132, -1.6271]], dtype=torch.float64)
[8]:
# Compute beta functions at observation points and corresponding derivatives with respect to deviation parameters

def fn(knobs):

    bxs = []
    bys = []

    m = torch.func.jacfwd(ring)(state, knobs)

    *_, w = twiss(m)
    _, bx, _, by = wolski_to_cs(w)

    for line in lines:
        w = propagate(w, torch.func.jacrev(line)(state, knobs))
        _, bx, _, by = wolski_to_cs(w)
        bxs.append(bx)
        bys.append(by)

    bxs = torch.stack(bxs)
    bys = torch.stack(bys)

    return bxs, bys

bx, by = fn(knobs)
dbxdk, dbydk = torch.func.jacrev(fn)(knobs)

print(bx)
print(dbxdk)
print()

print(by)
print(dbydk)
print()
tensor([18.6083, 18.6083, 18.6083, 18.6083], dtype=torch.float64)
tensor([[ 21.1373,  -1.5714,  21.1373,  -1.5714, 140.4889, -10.4446, 140.4889, -10.4446],
        [140.4889, -10.4446,  21.1373,  -1.5714,  21.1373,  -1.5714, 140.4889, -10.4446],
        [140.4889, -10.4446, 140.4889, -10.4446,  21.1373,  -1.5714,  21.1373,  -1.5714],
        [ 21.1373,  -1.5714, 140.4889, -10.4446, 140.4889, -10.4446,  21.1373,  -1.5714]], dtype=torch.float64)

tensor([6.3291, 6.3291, 6.3291, 6.3291], dtype=torch.float64)
tensor([[ 10.9592,  66.8187,  10.9592,  66.8187,  -5.0152, -30.5777,  -5.0152, -30.5777],
        [ -5.0152, -30.5777,  10.9592,  66.8187,  10.9592,  66.8187,  -5.0152, -30.5777],
        [ -5.0152, -30.5777,  -5.0152, -30.5777,  10.9592,  66.8187,  10.9592,  66.8187],
        [ 10.9592,  66.8187,  -5.0152, -30.5777,  -5.0152, -30.5777,  10.9592,  66.8187]], dtype=torch.float64)

Example-24: Module

[1]:
# Given an objective function or a model, it can be wrapped with torch.nn.Module
# This allows to use different optimization methods (torch.optim, pytorch-optimizer, ...)

# In this example chromaticity is optimized by
# 1) wrapping objective funtion (R^n x R^m x ... -> R) with a torch module (no data is passed to forward call)
# 2) wrapping chromaticity function with a torch module (plane index is passed as feature to forward call)

# In the first case, regular optimization is performed using all avaliable data

# For the second case, mini-batched optimization can be performed
# Planes (horozontal or vertical) are used as features
# Alternatively, location indices along the ring can be used as features (values of twiss parameters at location)
# Or location pairs (phase advance)
[2]:
# Import

import torch
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader
torch.set_printoptions(linewidth=128)

import matplotlib
from matplotlib import pyplot as plt
matplotlib.rcParams['text.usetex'] = True

from twiss import twiss

from ndmap.signature import chop
from ndmap.evaluate import evaluate
from ndmap.pfp import parametric_fixed_point

from model.library.drift import Drift
from model.library.quadrupole import Quadrupole
from model.library.sextupole import Sextupole
from model.library.dipole import Dipole
from model.library.line import Line

from model.command.wrapper import group
from model.command.wrapper import Wrapper
[3]:
# Define simple FODO based lattice using nested lines

DR = Drift('DR', 0.25)
BM = Dipole('BM', 3.50, torch.pi/4.0)

QF_A = Quadrupole('QF_A', 0.5, +0.20)
QD_A = Quadrupole('QD_A', 0.5, -0.19)
QF_B = Quadrupole('QF_B', 0.5, +0.20)
QD_B = Quadrupole('QD_B', 0.5, -0.19)
QF_C = Quadrupole('QF_C', 0.5, +0.20)
QD_C = Quadrupole('QD_C', 0.5, -0.19)
QF_D = Quadrupole('QF_D', 0.5, +0.20)
QD_D = Quadrupole('QD_D', 0.5, -0.19)

SF_A = Sextupole('SF_A', 0.25, 0.00)
SD_A = Sextupole('SD_A', 0.25, 0.00)
SF_B = Sextupole('SF_B', 0.25, 0.00)
SD_B = Sextupole('SD_B', 0.25, 0.00)
SF_C = Sextupole('SF_C', 0.25, 0.00)
SD_C = Sextupole('SD_C', 0.25, 0.00)
SF_D = Sextupole('SF_D', 0.25, 0.00)
SD_D = Sextupole('SD_D', 0.25, 0.00)

FODO_A = Line('FODO_A', [QF_A, DR, SF_A, DR, BM, DR, SD_A, DR, QD_A, QD_A, DR, SD_A, DR, BM, DR, SF_A, DR, QF_A], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_B = Line('FODO_B', [QF_B, DR, SF_B, DR, BM, DR, SD_B, DR, QD_B, QD_B, DR, SD_B, DR, BM, DR, SF_B, DR, QF_B], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_C = Line('FODO_C', [QF_C, DR, SF_C, DR, BM, DR, SD_C, DR, QD_C, QD_C, DR, SD_C, DR, BM, DR, SF_C, DR, QF_C], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_D = Line('FODO_D', [QF_D, DR, SF_D, DR, BM, DR, SD_D, DR, QD_D, QD_D, DR, SD_D, DR, BM, DR, SF_D, DR, QF_D], propagate=True, dp=0.0, exact=False, output=False, matrix=False)

RING = Line('RING', [FODO_A, FODO_B, FODO_C, FODO_D], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
[4]:
# Set parametric mapping

ring, *_ = group(RING, 'FODO_A', 'FODO_D', ('ms', ['Sextupole'], None, None), ('dp', None, None, None), root=True)

# Set deviation parameters

fp = torch.tensor(4*[0.0], dtype=torch.float64)
ms = torch.tensor(8*[0.0], dtype=torch.float64)
dp = torch.tensor([0.0], dtype=torch.float64)
[5]:
# Define parametric chomaticity function

# Compute parametric fixed point (first order dispersion)

pfp, *_ = parametric_fixed_point((0, 1), fp, [ms, dp], ring)
chop(pfp)

# Define ring around parametric fixed point

def mapping(state, ms, dp):
    return ring(state + evaluate(pfp, [ms, dp]), ms, dp) - evaluate(pfp, [ms, dp])

# Define tunes

def tune(ms, dp):
    matrix = torch.func.jacrev(mapping)(fp, ms, dp)
    tunes, *_ = twiss(matrix)
    return tunes

# Define chromaticity

def chromaticity(ms):
    return torch.func.jacrev(tune, 1)(ms, dp).squeeze()

# Compute natural chromaticity

print(chromaticity(ms))
tensor([-2.0649, -0.8260], dtype=torch.float64)
[6]:
# Chromaticity can be corrected in a single step

# Compute starting values

psix, psiy = chromaticity(ms)

# Set target values

psix_target = torch.tensor(5.0, dtype=torch.float64)
psiy_target = torch.tensor(5.0, dtype=torch.float64)

# Perform correction

dpsix = psix - psix_target
dpsiy = psiy - psiy_target

solution = - torch.linalg.pinv((torch.func.jacrev(chromaticity)(ms)).squeeze()) @ torch.stack([dpsix, dpsiy])
print(solution)

# Test solution

print(chromaticity(solution))
tensor([ 0.7439, -1.2084,  0.7439, -1.2084,  0.7439, -1.2084,  0.7439, -1.2084], dtype=torch.float64)
tensor([5.0000, 5.0000], dtype=torch.float64)
[7]:
# Optimization (wrapping objective funtion)

# Set model parameters
# Parameters are not cloned inside the module on initialization, values will change during optimization!

ms = torch.tensor(8*[0.0], dtype=torch.float64)

# Define scalar objective function

def objective(ms):
    psix, psiy = chromaticity(ms)
    return ((psix - psix_target)**2 + (psiy - psiy_target)**2).sqrt()

# Set model (forward returns evaluated objective)

model = Wrapper(objective, ms)

# Set optimizer

optimizer = torch.optim.Adam(model.parameters(), lr=1.0E-2)

# Perfom optimization

epochs = 256
for epoch in range(epochs):

    # Evaluate model
    error = model()

    # Compute derivatives
    error.backward()

    # Perform optimization step
    optimizer.step()

    # Set gradient to zero
    optimizer.zero_grad()

    # Verbose
    knobs, *_ = [*model.parameters()]
    print(error.detach(), (knobs.detach() - solution).norm())
tensor(9.1573, dtype=torch.float64) tensor(2.8105, dtype=torch.float64)
tensor(9.0611, dtype=torch.float64) tensor(2.7830, dtype=torch.float64)
tensor(8.9651, dtype=torch.float64) tensor(2.7555, dtype=torch.float64)
tensor(8.8693, dtype=torch.float64) tensor(2.7280, dtype=torch.float64)
tensor(8.7737, dtype=torch.float64) tensor(2.7006, dtype=torch.float64)
tensor(8.6784, dtype=torch.float64) tensor(2.6732, dtype=torch.float64)
tensor(8.5833, dtype=torch.float64) tensor(2.6458, dtype=torch.float64)
tensor(8.4884, dtype=torch.float64) tensor(2.6184, dtype=torch.float64)
tensor(8.3938, dtype=torch.float64) tensor(2.5910, dtype=torch.float64)
tensor(8.2995, dtype=torch.float64) tensor(2.5636, dtype=torch.float64)
tensor(8.2054, dtype=torch.float64) tensor(2.5363, dtype=torch.float64)
tensor(8.1116, dtype=torch.float64) tensor(2.5090, dtype=torch.float64)
tensor(8.0181, dtype=torch.float64) tensor(2.4817, dtype=torch.float64)
tensor(7.9249, dtype=torch.float64) tensor(2.4544, dtype=torch.float64)
tensor(7.8320, dtype=torch.float64) tensor(2.4271, dtype=torch.float64)
tensor(7.7394, dtype=torch.float64) tensor(2.3999, dtype=torch.float64)
tensor(7.6471, dtype=torch.float64) tensor(2.3727, dtype=torch.float64)
tensor(7.5552, dtype=torch.float64) tensor(2.3455, dtype=torch.float64)
tensor(7.4636, dtype=torch.float64) tensor(2.3183, dtype=torch.float64)
tensor(7.3724, dtype=torch.float64) tensor(2.2912, dtype=torch.float64)
tensor(7.2815, dtype=torch.float64) tensor(2.2641, dtype=torch.float64)
tensor(7.1910, dtype=torch.float64) tensor(2.2370, dtype=torch.float64)
tensor(7.1008, dtype=torch.float64) tensor(2.2099, dtype=torch.float64)
tensor(7.0110, dtype=torch.float64) tensor(2.1829, dtype=torch.float64)
tensor(6.9216, dtype=torch.float64) tensor(2.1560, dtype=torch.float64)
tensor(6.8326, dtype=torch.float64) tensor(2.1290, dtype=torch.float64)
tensor(6.7440, dtype=torch.float64) tensor(2.1021, dtype=torch.float64)
tensor(6.6557, dtype=torch.float64) tensor(2.0752, dtype=torch.float64)
tensor(6.5679, dtype=torch.float64) tensor(2.0484, dtype=torch.float64)
tensor(6.4804, dtype=torch.float64) tensor(2.0216, dtype=torch.float64)
tensor(6.3933, dtype=torch.float64) tensor(1.9948, dtype=torch.float64)
tensor(6.3066, dtype=torch.float64) tensor(1.9681, dtype=torch.float64)
tensor(6.2204, dtype=torch.float64) tensor(1.9415, dtype=torch.float64)
tensor(6.1345, dtype=torch.float64) tensor(1.9148, dtype=torch.float64)
tensor(6.0489, dtype=torch.float64) tensor(1.8883, dtype=torch.float64)
tensor(5.9638, dtype=torch.float64) tensor(1.8617, dtype=torch.float64)
tensor(5.8791, dtype=torch.float64) tensor(1.8353, dtype=torch.float64)
tensor(5.7947, dtype=torch.float64) tensor(1.8089, dtype=torch.float64)
tensor(5.7107, dtype=torch.float64) tensor(1.7825, dtype=torch.float64)
tensor(5.6271, dtype=torch.float64) tensor(1.7562, dtype=torch.float64)
tensor(5.5438, dtype=torch.float64) tensor(1.7299, dtype=torch.float64)
tensor(5.4609, dtype=torch.float64) tensor(1.7037, dtype=torch.float64)
tensor(5.3783, dtype=torch.float64) tensor(1.6776, dtype=torch.float64)
tensor(5.2961, dtype=torch.float64) tensor(1.6515, dtype=torch.float64)
tensor(5.2142, dtype=torch.float64) tensor(1.6255, dtype=torch.float64)
tensor(5.1326, dtype=torch.float64) tensor(1.5996, dtype=torch.float64)
tensor(5.0513, dtype=torch.float64) tensor(1.5737, dtype=torch.float64)
tensor(4.9702, dtype=torch.float64) tensor(1.5479, dtype=torch.float64)
tensor(4.8895, dtype=torch.float64) tensor(1.5221, dtype=torch.float64)
tensor(4.8091, dtype=torch.float64) tensor(1.4964, dtype=torch.float64)
tensor(4.7289, dtype=torch.float64) tensor(1.4708, dtype=torch.float64)
tensor(4.6489, dtype=torch.float64) tensor(1.4453, dtype=torch.float64)
tensor(4.5692, dtype=torch.float64) tensor(1.4198, dtype=torch.float64)
tensor(4.4897, dtype=torch.float64) tensor(1.3944, dtype=torch.float64)
tensor(4.4105, dtype=torch.float64) tensor(1.3690, dtype=torch.float64)
tensor(4.3314, dtype=torch.float64) tensor(1.3438, dtype=torch.float64)
tensor(4.2525, dtype=torch.float64) tensor(1.3186, dtype=torch.float64)
tensor(4.1738, dtype=torch.float64) tensor(1.2934, dtype=torch.float64)
tensor(4.0953, dtype=torch.float64) tensor(1.2683, dtype=torch.float64)
tensor(4.0169, dtype=torch.float64) tensor(1.2433, dtype=torch.float64)
tensor(3.9387, dtype=torch.float64) tensor(1.2184, dtype=torch.float64)
tensor(3.8606, dtype=torch.float64) tensor(1.1935, dtype=torch.float64)
tensor(3.7827, dtype=torch.float64) tensor(1.1687, dtype=torch.float64)
tensor(3.7049, dtype=torch.float64) tensor(1.1440, dtype=torch.float64)
tensor(3.6272, dtype=torch.float64) tensor(1.1193, dtype=torch.float64)
tensor(3.5496, dtype=torch.float64) tensor(1.0947, dtype=torch.float64)
tensor(3.4722, dtype=torch.float64) tensor(1.0702, dtype=torch.float64)
tensor(3.3948, dtype=torch.float64) tensor(1.0457, dtype=torch.float64)
tensor(3.3176, dtype=torch.float64) tensor(1.0212, dtype=torch.float64)
tensor(3.2404, dtype=torch.float64) tensor(0.9968, dtype=torch.float64)
tensor(3.1634, dtype=torch.float64) tensor(0.9725, dtype=torch.float64)
tensor(3.0864, dtype=torch.float64) tensor(0.9482, dtype=torch.float64)
tensor(3.0096, dtype=torch.float64) tensor(0.9240, dtype=torch.float64)
tensor(2.9328, dtype=torch.float64) tensor(0.8998, dtype=torch.float64)
tensor(2.8561, dtype=torch.float64) tensor(0.8757, dtype=torch.float64)
tensor(2.7796, dtype=torch.float64) tensor(0.8516, dtype=torch.float64)
tensor(2.7031, dtype=torch.float64) tensor(0.8276, dtype=torch.float64)
tensor(2.6267, dtype=torch.float64) tensor(0.8035, dtype=torch.float64)
tensor(2.5504, dtype=torch.float64) tensor(0.7796, dtype=torch.float64)
tensor(2.4742, dtype=torch.float64) tensor(0.7556, dtype=torch.float64)
tensor(2.3980, dtype=torch.float64) tensor(0.7317, dtype=torch.float64)
tensor(2.3220, dtype=torch.float64) tensor(0.7078, dtype=torch.float64)
tensor(2.2460, dtype=torch.float64) tensor(0.6839, dtype=torch.float64)
tensor(2.1702, dtype=torch.float64) tensor(0.6601, dtype=torch.float64)
tensor(2.0944, dtype=torch.float64) tensor(0.6363, dtype=torch.float64)
tensor(2.0186, dtype=torch.float64) tensor(0.6125, dtype=torch.float64)
tensor(1.9430, dtype=torch.float64) tensor(0.5887, dtype=torch.float64)
tensor(1.8674, dtype=torch.float64) tensor(0.5649, dtype=torch.float64)
tensor(1.7919, dtype=torch.float64) tensor(0.5412, dtype=torch.float64)
tensor(1.7165, dtype=torch.float64) tensor(0.5174, dtype=torch.float64)
tensor(1.6411, dtype=torch.float64) tensor(0.4937, dtype=torch.float64)
tensor(1.5658, dtype=torch.float64) tensor(0.4699, dtype=torch.float64)
tensor(1.4906, dtype=torch.float64) tensor(0.4462, dtype=torch.float64)
tensor(1.4154, dtype=torch.float64) tensor(0.4225, dtype=torch.float64)
tensor(1.3402, dtype=torch.float64) tensor(0.3988, dtype=torch.float64)
tensor(1.2652, dtype=torch.float64) tensor(0.3751, dtype=torch.float64)
tensor(1.1901, dtype=torch.float64) tensor(0.3514, dtype=torch.float64)
tensor(1.1152, dtype=torch.float64) tensor(0.3278, dtype=torch.float64)
tensor(1.0402, dtype=torch.float64) tensor(0.3041, dtype=torch.float64)
tensor(0.9654, dtype=torch.float64) tensor(0.2805, dtype=torch.float64)
tensor(0.8906, dtype=torch.float64) tensor(0.2569, dtype=torch.float64)
tensor(0.8158, dtype=torch.float64) tensor(0.2333, dtype=torch.float64)
tensor(0.7412, dtype=torch.float64) tensor(0.2098, dtype=torch.float64)
tensor(0.6665, dtype=torch.float64) tensor(0.1862, dtype=torch.float64)
tensor(0.5920, dtype=torch.float64) tensor(0.1627, dtype=torch.float64)
tensor(0.5174, dtype=torch.float64) tensor(0.1393, dtype=torch.float64)
tensor(0.4430, dtype=torch.float64) tensor(0.1159, dtype=torch.float64)
tensor(0.3685, dtype=torch.float64) tensor(0.0925, dtype=torch.float64)
tensor(0.2941, dtype=torch.float64) tensor(0.0691, dtype=torch.float64)
tensor(0.2198, dtype=torch.float64) tensor(0.0458, dtype=torch.float64)
tensor(0.1455, dtype=torch.float64) tensor(0.0225, dtype=torch.float64)
tensor(0.0713, dtype=torch.float64) tensor(0.0009, dtype=torch.float64)
tensor(0.0034, dtype=torch.float64) tensor(0.0210, dtype=torch.float64)
tensor(0.0681, dtype=torch.float64) tensor(0.0360, dtype=torch.float64)
tensor(0.1144, dtype=torch.float64) tensor(0.0467, dtype=torch.float64)
tensor(0.1474, dtype=torch.float64) tensor(0.0539, dtype=torch.float64)
tensor(0.1708, dtype=torch.float64) tensor(0.0581, dtype=torch.float64)
tensor(0.1860, dtype=torch.float64) tensor(0.0598, dtype=torch.float64)
tensor(0.1928, dtype=torch.float64) tensor(0.0591, dtype=torch.float64)
tensor(0.1912, dtype=torch.float64) tensor(0.0563, dtype=torch.float64)
tensor(0.1816, dtype=torch.float64) tensor(0.0516, dtype=torch.float64)
tensor(0.1652, dtype=torch.float64) tensor(0.0453, dtype=torch.float64)
tensor(0.1436, dtype=torch.float64) tensor(0.0373, dtype=torch.float64)
tensor(0.1178, dtype=torch.float64) tensor(0.0277, dtype=torch.float64)
tensor(0.0874, dtype=torch.float64) tensor(0.0163, dtype=torch.float64)
tensor(0.0514, dtype=torch.float64) tensor(0.0036, dtype=torch.float64)
tensor(0.0133, dtype=torch.float64) tensor(0.0093, dtype=torch.float64)
tensor(0.0305, dtype=torch.float64) tensor(0.0187, dtype=torch.float64)
tensor(0.0602, dtype=torch.float64) tensor(0.0250, dtype=torch.float64)
tensor(0.0796, dtype=torch.float64) tensor(0.0285, dtype=torch.float64)
tensor(0.0901, dtype=torch.float64) tensor(0.0292, dtype=torch.float64)
tensor(0.0923, dtype=torch.float64) tensor(0.0274, dtype=torch.float64)
tensor(0.0868, dtype=torch.float64) tensor(0.0234, dtype=torch.float64)
tensor(0.0744, dtype=torch.float64) tensor(0.0175, dtype=torch.float64)
tensor(0.0563, dtype=torch.float64) tensor(0.0101, dtype=torch.float64)
tensor(0.0326, dtype=torch.float64) tensor(0.0014, dtype=torch.float64)
tensor(0.0048, dtype=torch.float64) tensor(0.0101, dtype=torch.float64)
tensor(0.0362, dtype=torch.float64) tensor(0.0164, dtype=torch.float64)
tensor(0.0519, dtype=torch.float64) tensor(0.0199, dtype=torch.float64)
tensor(0.0647, dtype=torch.float64) tensor(0.0213, dtype=torch.float64)
tensor(0.0721, dtype=torch.float64) tensor(0.0205, dtype=torch.float64)
tensor(0.0681, dtype=torch.float64) tensor(0.0179, dtype=torch.float64)
tensor(0.0566, dtype=torch.float64) tensor(0.0135, dtype=torch.float64)
tensor(0.0438, dtype=torch.float64) tensor(0.0067, dtype=torch.float64)
tensor(0.0224, dtype=torch.float64) tensor(0.0036, dtype=torch.float64)
tensor(0.0187, dtype=torch.float64) tensor(0.0082, dtype=torch.float64)
tensor(0.0269, dtype=torch.float64) tensor(0.0118, dtype=torch.float64)
tensor(0.0412, dtype=torch.float64) tensor(0.0129, dtype=torch.float64)
tensor(0.0420, dtype=torch.float64) tensor(0.0122, dtype=torch.float64)
tensor(0.0388, dtype=torch.float64) tensor(0.0090, dtype=torch.float64)
tensor(0.0296, dtype=torch.float64) tensor(0.0034, dtype=torch.float64)
tensor(0.0122, dtype=torch.float64) tensor(0.0032, dtype=torch.float64)
tensor(0.0101, dtype=torch.float64) tensor(0.0068, dtype=torch.float64)
tensor(0.0220, dtype=torch.float64) tensor(0.0079, dtype=torch.float64)
tensor(0.0250, dtype=torch.float64) tensor(0.0066, dtype=torch.float64)
tensor(0.0208, dtype=torch.float64) tensor(0.0030, dtype=torch.float64)
tensor(0.0100, dtype=torch.float64) tensor(0.0023, dtype=torch.float64)
tensor(0.0092, dtype=torch.float64) tensor(0.0063, dtype=torch.float64)
tensor(0.0233, dtype=torch.float64) tensor(0.0064, dtype=torch.float64)
tensor(0.0202, dtype=torch.float64) tensor(0.0052, dtype=torch.float64)
tensor(0.0222, dtype=torch.float64) tensor(0.0022, dtype=torch.float64)
tensor(0.0071, dtype=torch.float64) tensor(0.0026, dtype=torch.float64)
tensor(0.0109, dtype=torch.float64) tensor(0.0063, dtype=torch.float64)
tensor(0.0244, dtype=torch.float64) tensor(0.0062, dtype=torch.float64)
tensor(0.0201, dtype=torch.float64) tensor(0.0050, dtype=torch.float64)
tensor(0.0235, dtype=torch.float64) tensor(0.0020, dtype=torch.float64)
tensor(0.0084, dtype=torch.float64) tensor(0.0058, dtype=torch.float64)
tensor(0.0370, dtype=torch.float64) tensor(0.0067, dtype=torch.float64)
tensor(0.0358, dtype=torch.float64) tensor(0.0074, dtype=torch.float64)
tensor(0.0242, dtype=torch.float64) tensor(0.0087, dtype=torch.float64)
tensor(0.0366, dtype=torch.float64) tensor(0.0057, dtype=torch.float64)
tensor(0.0215, dtype=torch.float64) tensor(0.0043, dtype=torch.float64)
tensor(0.0269, dtype=torch.float64) tensor(0.0042, dtype=torch.float64)
tensor(0.0241, dtype=torch.float64) tensor(0.0046, dtype=torch.float64)
tensor(0.0212, dtype=torch.float64) tensor(0.0058, dtype=torch.float64)
tensor(0.0275, dtype=torch.float64) tensor(0.0046, dtype=torch.float64)
tensor(0.0146, dtype=torch.float64) tensor(0.0036, dtype=torch.float64)
tensor(0.0189, dtype=torch.float64) tensor(0.0014, dtype=torch.float64)
tensor(0.0056, dtype=torch.float64) tensor(0.0025, dtype=torch.float64)
tensor(0.0087, dtype=torch.float64) tensor(0.0019, dtype=torch.float64)
tensor(0.0073, dtype=torch.float64) tensor(0.0022, dtype=torch.float64)
tensor(0.0128, dtype=torch.float64) tensor(0.0025, dtype=torch.float64)
tensor(0.0102, dtype=torch.float64) tensor(0.0023, dtype=torch.float64)
tensor(0.0072, dtype=torch.float64) tensor(0.0012, dtype=torch.float64)
tensor(0.0076, dtype=torch.float64) tensor(0.0040, dtype=torch.float64)
tensor(0.0213, dtype=torch.float64) tensor(0.0039, dtype=torch.float64)
tensor(0.0158, dtype=torch.float64) tensor(0.0042, dtype=torch.float64)
tensor(0.0231, dtype=torch.float64) tensor(0.0033, dtype=torch.float64)
tensor(0.0202, dtype=torch.float64) tensor(0.0023, dtype=torch.float64)
tensor(0.0146, dtype=torch.float64) tensor(0.0023, dtype=torch.float64)
tensor(0.0128, dtype=torch.float64) tensor(0.0039, dtype=torch.float64)
tensor(0.0218, dtype=torch.float64) tensor(0.0035, dtype=torch.float64)
tensor(0.0197, dtype=torch.float64) tensor(0.0023, dtype=torch.float64)
tensor(0.0139, dtype=torch.float64) tensor(0.0018, dtype=torch.float64)
tensor(0.0115, dtype=torch.float64) tensor(0.0036, dtype=torch.float64)
tensor(0.0222, dtype=torch.float64) tensor(0.0036, dtype=torch.float64)
tensor(0.0215, dtype=torch.float64) tensor(0.0022, dtype=torch.float64)
tensor(0.0111, dtype=torch.float64) tensor(0.0016, dtype=torch.float64)
tensor(0.0100, dtype=torch.float64) tensor(0.0033, dtype=torch.float64)
tensor(0.0204, dtype=torch.float64) tensor(0.0032, dtype=torch.float64)
tensor(0.0180, dtype=torch.float64) tensor(0.0028, dtype=torch.float64)
tensor(0.0152, dtype=torch.float64) tensor(0.0025, dtype=torch.float64)
tensor(0.0149, dtype=torch.float64) tensor(0.0021, dtype=torch.float64)
tensor(0.0137, dtype=torch.float64) tensor(0.0019, dtype=torch.float64)
tensor(0.0103, dtype=torch.float64) tensor(0.0036, dtype=torch.float64)
tensor(0.0216, dtype=torch.float64) tensor(0.0034, dtype=torch.float64)
tensor(0.0209, dtype=torch.float64) tensor(0.0014, dtype=torch.float64)
tensor(0.0084, dtype=torch.float64) tensor(0.0010, dtype=torch.float64)
tensor(0.0063, dtype=torch.float64) tensor(0.0038, dtype=torch.float64)
tensor(0.0236, dtype=torch.float64) tensor(0.0036, dtype=torch.float64)
tensor(0.0222, dtype=torch.float64) tensor(0.0013, dtype=torch.float64)
tensor(0.0074, dtype=torch.float64) tensor(0.0010, dtype=torch.float64)
tensor(0.0065, dtype=torch.float64) tensor(0.0034, dtype=torch.float64)
tensor(0.0216, dtype=torch.float64) tensor(0.0032, dtype=torch.float64)
tensor(0.0196, dtype=torch.float64) tensor(0.0018, dtype=torch.float64)
tensor(0.0101, dtype=torch.float64) tensor(0.0015, dtype=torch.float64)
tensor(0.0097, dtype=torch.float64) tensor(0.0028, dtype=torch.float64)
tensor(0.0174, dtype=torch.float64) tensor(0.0026, dtype=torch.float64)
tensor(0.0151, dtype=torch.float64) tensor(0.0024, dtype=torch.float64)
tensor(0.0141, dtype=torch.float64) tensor(0.0021, dtype=torch.float64)
tensor(0.0134, dtype=torch.float64) tensor(0.0021, dtype=torch.float64)
tensor(0.0134, dtype=torch.float64) tensor(0.0019, dtype=torch.float64)
tensor(0.0112, dtype=torch.float64) tensor(0.0028, dtype=torch.float64)
tensor(0.0171, dtype=torch.float64) tensor(0.0026, dtype=torch.float64)
tensor(0.0161, dtype=torch.float64) tensor(0.0017, dtype=torch.float64)
tensor(0.0108, dtype=torch.float64) tensor(0.0015, dtype=torch.float64)
tensor(0.0089, dtype=torch.float64) tensor(0.0029, dtype=torch.float64)
tensor(0.0186, dtype=torch.float64) tensor(0.0027, dtype=torch.float64)
tensor(0.0174, dtype=torch.float64) tensor(0.0015, dtype=torch.float64)
tensor(0.0093, dtype=torch.float64) tensor(0.0012, dtype=torch.float64)
tensor(0.0076, dtype=torch.float64) tensor(0.0030, dtype=torch.float64)
tensor(0.0192, dtype=torch.float64) tensor(0.0029, dtype=torch.float64)
tensor(0.0180, dtype=torch.float64) tensor(0.0013, dtype=torch.float64)
tensor(0.0084, dtype=torch.float64) tensor(0.0011, dtype=torch.float64)
tensor(0.0069, dtype=torch.float64) tensor(0.0031, dtype=torch.float64)
tensor(0.0195, dtype=torch.float64) tensor(0.0028, dtype=torch.float64)
tensor(0.0182, dtype=torch.float64) tensor(0.0013, dtype=torch.float64)
tensor(0.0079, dtype=torch.float64) tensor(0.0010, dtype=torch.float64)
tensor(0.0065, dtype=torch.float64) tensor(0.0031, dtype=torch.float64)
tensor(0.0193, dtype=torch.float64) tensor(0.0029, dtype=torch.float64)
tensor(0.0180, dtype=torch.float64) tensor(0.0013, dtype=torch.float64)
tensor(0.0079, dtype=torch.float64) tensor(0.0011, dtype=torch.float64)
tensor(0.0068, dtype=torch.float64) tensor(0.0029, dtype=torch.float64)
tensor(0.0186, dtype=torch.float64) tensor(0.0027, dtype=torch.float64)
tensor(0.0171, dtype=torch.float64) tensor(0.0015, dtype=torch.float64)
tensor(0.0087, dtype=torch.float64) tensor(0.0012, dtype=torch.float64)
tensor(0.0077, dtype=torch.float64) tensor(0.0027, dtype=torch.float64)
tensor(0.0172, dtype=torch.float64) tensor(0.0026, dtype=torch.float64)
tensor(0.0156, dtype=torch.float64) tensor(0.0016, dtype=torch.float64)
tensor(0.0101, dtype=torch.float64) tensor(0.0014, dtype=torch.float64)
tensor(0.0092, dtype=torch.float64) tensor(0.0024, dtype=torch.float64)
tensor(0.0153, dtype=torch.float64) tensor(0.0022, dtype=torch.float64)
tensor(0.0137, dtype=torch.float64) tensor(0.0020, dtype=torch.float64)
tensor(0.0117, dtype=torch.float64) tensor(0.0017, dtype=torch.float64)
tensor(0.0108, dtype=torch.float64) tensor(0.0022, dtype=torch.float64)
tensor(0.0136, dtype=torch.float64) tensor(0.0020, dtype=torch.float64)
[8]:
# Optimization (wrapping chromaticity function)

# Set model parameters
# Parameters are not cloned inside the module on initialization, values will change during optimization!

ms = torch.tensor(8*[0.0], dtype=torch.float64)

# Set features and labels
# X selects the plane (horizontal or vertical chomaticity)
# y is corresponding target chromaticity value for selected plane

X = torch.tensor([[0], [1]])
y = torch.stack([psix_target, psiy_target])

# Set dataset
# Note, batch size is one, technicaly this is not a mini-batch optimization

batch_size = 1
dataset = TensorDataset(X.clone(), y.clone())

# Set data loader

dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# Set objective (return horizontal or vertical chomaticity)

def objective(x, ms):
    return chromaticity(ms)[x].squeeze()

# Set model (forward returns evaluated objective)

model = Wrapper(objective, ms)

# Set optimizer

optimizer = torch.optim.Adam(model.parameters(), lr=1.0E-2)

# Set loss funtion

lf = torch.nn.MSELoss()

# Perfom optimization

epochs = 256
for epoch in range(epochs):

    # Loop over batches of data
    for batch, (X, y) in enumerate(dataloader):

        # Evaluate model
        y_hat = model(X)

        # Evaluate loss function
        error = lf(y_hat, y.squeeze())

        # Compute derivatives
        error.backward()

        # Perform optimization step
        optimizer.step()

        # Set gradient to zero
        optimizer.zero_grad()

    # Verbose
    knobs, *_ = [*model.parameters()]
    print(error.detach(), (knobs.detach() - solution).norm())
tensor(35.4295, dtype=torch.float64) tensor(2.8297, dtype=torch.float64)
tensor(47.7980, dtype=torch.float64) tensor(2.8006, dtype=torch.float64)
tensor(33.9269, dtype=torch.float64) tensor(2.7718, dtype=torch.float64)
tensor(45.6336, dtype=torch.float64) tensor(2.7423, dtype=torch.float64)
tensor(32.6024, dtype=torch.float64) tensor(2.7134, dtype=torch.float64)
tensor(43.5373, dtype=torch.float64) tensor(2.6840, dtype=torch.float64)
tensor(42.6609, dtype=torch.float64) tensor(2.6547, dtype=torch.float64)
tensor(41.7582, dtype=torch.float64) tensor(2.6255, dtype=torch.float64)
tensor(29.9431, dtype=torch.float64) tensor(2.5970, dtype=torch.float64)
tensor(29.4738, dtype=torch.float64) tensor(2.5684, dtype=torch.float64)
tensor(28.9800, dtype=torch.float64) tensor(2.5399, dtype=torch.float64)
tensor(37.5585, dtype=torch.float64) tensor(2.5109, dtype=torch.float64)
tensor(36.7554, dtype=torch.float64) tensor(2.4820, dtype=torch.float64)
tensor(27.1260, dtype=torch.float64) tensor(2.4540, dtype=torch.float64)
tensor(26.6535, dtype=torch.float64) tensor(2.4259, dtype=torch.float64)
tensor(26.1669, dtype=torch.float64) tensor(2.3979, dtype=torch.float64)
tensor(33.0411, dtype=torch.float64) tensor(2.3695, dtype=torch.float64)
tensor(32.3221, dtype=torch.float64) tensor(2.3413, dtype=torch.float64)
tensor(31.5935, dtype=torch.float64) tensor(2.3133, dtype=torch.float64)
tensor(30.8600, dtype=torch.float64) tensor(2.2855, dtype=torch.float64)
tensor(23.2333, dtype=torch.float64) tensor(2.2585, dtype=torch.float64)
tensor(22.8195, dtype=torch.float64) tensor(2.2316, dtype=torch.float64)
tensor(22.3918, dtype=torch.float64) tensor(2.2048, dtype=torch.float64)
tensor(27.5644, dtype=torch.float64) tensor(2.1774, dtype=torch.float64)
tensor(21.3774, dtype=torch.float64) tensor(2.1510, dtype=torch.float64)
tensor(20.9557, dtype=torch.float64) tensor(2.1246, dtype=torch.float64)
tensor(25.3956, dtype=torch.float64) tensor(2.0977, dtype=torch.float64)
tensor(24.8144, dtype=torch.float64) tensor(2.0711, dtype=torch.float64)
tensor(24.2275, dtype=torch.float64) tensor(2.0447, dtype=torch.float64)
tensor(18.9254, dtype=torch.float64) tensor(2.0192, dtype=torch.float64)
tensor(22.9146, dtype=torch.float64) tensor(1.9932, dtype=torch.float64)
tensor(18.0572, dtype=torch.float64) tensor(1.9680, dtype=torch.float64)
tensor(17.6949, dtype=torch.float64) tensor(1.9429, dtype=torch.float64)
tensor(21.0109, dtype=torch.float64) tensor(1.9173, dtype=torch.float64)
tensor(20.5056, dtype=torch.float64) tensor(1.8920, dtype=torch.float64)
tensor(16.3728, dtype=torch.float64) tensor(1.8675, dtype=torch.float64)
tensor(16.0336, dtype=torch.float64) tensor(1.8432, dtype=torch.float64)
tensor(18.7730, dtype=torch.float64) tensor(1.8183, dtype=torch.float64)
tensor(15.2336, dtype=torch.float64) tensor(1.7943, dtype=torch.float64)
tensor(14.9012, dtype=torch.float64) tensor(1.7705, dtype=torch.float64)
tensor(14.5642, dtype=torch.float64) tensor(1.7468, dtype=torch.float64)
tensor(14.2246, dtype=torch.float64) tensor(1.7232, dtype=torch.float64)
tensor(13.8840, dtype=torch.float64) tensor(1.6997, dtype=torch.float64)
tensor(15.6818, dtype=torch.float64) tensor(1.6758, dtype=torch.float64)
tensor(13.1131, dtype=torch.float64) tensor(1.6528, dtype=torch.float64)
tensor(14.8403, dtype=torch.float64) tensor(1.6294, dtype=torch.float64)
tensor(14.4807, dtype=torch.float64) tensor(1.6062, dtype=torch.float64)
tensor(12.0058, dtype=torch.float64) tensor(1.5839, dtype=torch.float64)
tensor(11.7240, dtype=torch.float64) tensor(1.5617, dtype=torch.float64)
tensor(13.2212, dtype=torch.float64) tensor(1.5391, dtype=torch.float64)
tensor(12.8885, dtype=torch.float64) tensor(1.5167, dtype=torch.float64)
tensor(12.5528, dtype=torch.float64) tensor(1.4946, dtype=torch.float64)
tensor(12.2161, dtype=torch.float64) tensor(1.4727, dtype=torch.float64)
tensor(10.0792, dtype=torch.float64) tensor(1.4517, dtype=torch.float64)
tensor(9.8487, dtype=torch.float64) tensor(1.4307, dtype=torch.float64)
tensor(9.6134, dtype=torch.float64) tensor(1.4099, dtype=torch.float64)
tensor(10.7013, dtype=torch.float64) tensor(1.3887, dtype=torch.float64)
tensor(9.0662, dtype=torch.float64) tensor(1.3683, dtype=torch.float64)
tensor(8.8395, dtype=torch.float64) tensor(1.3480, dtype=torch.float64)
tensor(9.7233, dtype=torch.float64) tensor(1.3273, dtype=torch.float64)
tensor(9.4668, dtype=torch.float64) tensor(1.3069, dtype=torch.float64)
tensor(8.0403, dtype=torch.float64) tensor(1.2873, dtype=torch.float64)
tensor(8.8849, dtype=torch.float64) tensor(1.2672, dtype=torch.float64)
tensor(7.5717, dtype=torch.float64) tensor(1.2480, dtype=torch.float64)
tensor(7.3774, dtype=torch.float64) tensor(1.2289, dtype=torch.float64)
tensor(7.1814, dtype=torch.float64) tensor(1.2099, dtype=torch.float64)
tensor(7.7638, dtype=torch.float64) tensor(1.1906, dtype=torch.float64)
tensor(7.5544, dtype=torch.float64) tensor(1.1715, dtype=torch.float64)
tensor(6.4981, dtype=torch.float64) tensor(1.1532, dtype=torch.float64)
tensor(6.3243, dtype=torch.float64) tensor(1.1350, dtype=torch.float64)
tensor(6.8253, dtype=torch.float64) tensor(1.1165, dtype=torch.float64)
tensor(5.9263, dtype=torch.float64) tensor(1.0987, dtype=torch.float64)
tensor(5.7620, dtype=torch.float64) tensor(1.0811, dtype=torch.float64)
tensor(5.5974, dtype=torch.float64) tensor(1.0636, dtype=torch.float64)
tensor(5.9471, dtype=torch.float64) tensor(1.0458, dtype=torch.float64)
tensor(5.2268, dtype=torch.float64) tensor(1.0288, dtype=torch.float64)
tensor(5.0744, dtype=torch.float64) tensor(1.0119, dtype=torch.float64)
tensor(5.3731, dtype=torch.float64) tensor(0.9946, dtype=torch.float64)
tensor(4.7320, dtype=torch.float64) tensor(0.9782, dtype=torch.float64)
tensor(4.5913, dtype=torch.float64) tensor(0.9618, dtype=torch.float64)
tensor(4.8476, dtype=torch.float64) tensor(0.9452, dtype=torch.float64)
tensor(4.2758, dtype=torch.float64) tensor(0.9292, dtype=torch.float64)
tensor(4.5347, dtype=torch.float64) tensor(0.9130, dtype=torch.float64)
tensor(3.9830, dtype=torch.float64) tensor(0.8975, dtype=torch.float64)
tensor(4.2354, dtype=torch.float64) tensor(0.8816, dtype=torch.float64)
tensor(4.1097, dtype=torch.float64) tensor(0.8661, dtype=torch.float64)
tensor(3.9835, dtype=torch.float64) tensor(0.8507, dtype=torch.float64)
tensor(3.4304, dtype=torch.float64) tensor(0.8360, dtype=torch.float64)
tensor(3.3309, dtype=torch.float64) tensor(0.8215, dtype=torch.float64)
tensor(3.5530, dtype=torch.float64) tensor(0.8066, dtype=torch.float64)
tensor(3.4421, dtype=torch.float64) tensor(0.7920, dtype=torch.float64)
tensor(3.3313, dtype=torch.float64) tensor(0.7777, dtype=torch.float64)
tensor(3.2212, dtype=torch.float64) tensor(0.7635, dtype=torch.float64)
tensor(2.7570, dtype=torch.float64) tensor(0.7499, dtype=torch.float64)
tensor(2.9791, dtype=torch.float64) tensor(0.7361, dtype=torch.float64)
tensor(2.8788, dtype=torch.float64) tensor(0.7225, dtype=torch.float64)
tensor(2.4735, dtype=torch.float64) tensor(0.7095, dtype=torch.float64)
tensor(2.4011, dtype=torch.float64) tensor(0.6966, dtype=torch.float64)
tensor(2.3278, dtype=torch.float64) tensor(0.6838, dtype=torch.float64)
tensor(2.2539, dtype=torch.float64) tensor(0.6711, dtype=torch.float64)
tensor(2.3372, dtype=torch.float64) tensor(0.6583, dtype=torch.float64)
tensor(2.0867, dtype=torch.float64) tensor(0.6460, dtype=torch.float64)
tensor(2.1669, dtype=torch.float64) tensor(0.6335, dtype=torch.float64)
tensor(2.0960, dtype=torch.float64) tensor(0.6212, dtype=torch.float64)
tensor(1.8491, dtype=torch.float64) tensor(0.6095, dtype=torch.float64)
tensor(1.9372, dtype=torch.float64) tensor(0.5975, dtype=torch.float64)
tensor(1.7122, dtype=torch.float64) tensor(0.5861, dtype=torch.float64)
tensor(1.7894, dtype=torch.float64) tensor(0.5745, dtype=torch.float64)
tensor(1.5839, dtype=torch.float64) tensor(0.5634, dtype=torch.float64)
tensor(1.5311, dtype=torch.float64) tensor(0.5525, dtype=torch.float64)
tensor(1.5798, dtype=torch.float64) tensor(0.5413, dtype=torch.float64)
tensor(1.4119, dtype=torch.float64) tensor(0.5307, dtype=torch.float64)
tensor(1.3630, dtype=torch.float64) tensor(0.5202, dtype=torch.float64)
tensor(1.3145, dtype=torch.float64) tensor(0.5099, dtype=torch.float64)
tensor(1.2665, dtype=torch.float64) tensor(0.4997, dtype=torch.float64)
tensor(1.2191, dtype=torch.float64) tensor(0.4896, dtype=torch.float64)
tensor(1.1727, dtype=torch.float64) tensor(0.4797, dtype=torch.float64)
tensor(1.1272, dtype=torch.float64) tensor(0.4699, dtype=torch.float64)
tensor(1.0828, dtype=torch.float64) tensor(0.4602, dtype=torch.float64)
tensor(1.0396, dtype=torch.float64) tensor(0.4507, dtype=torch.float64)
tensor(1.0320, dtype=torch.float64) tensor(0.4411, dtype=torch.float64)
tensor(0.9468, dtype=torch.float64) tensor(0.4319, dtype=torch.float64)
tensor(0.9092, dtype=torch.float64) tensor(0.4229, dtype=torch.float64)
tensor(0.9150, dtype=torch.float64) tensor(0.4137, dtype=torch.float64)
tensor(0.8847, dtype=torch.float64) tensor(0.4048, dtype=torch.float64)
tensor(0.8541, dtype=torch.float64) tensor(0.3960, dtype=torch.float64)
tensor(0.7488, dtype=torch.float64) tensor(0.3876, dtype=torch.float64)
tensor(0.7846, dtype=torch.float64) tensor(0.3792, dtype=torch.float64)
tensor(0.7557, dtype=torch.float64) tensor(0.3709, dtype=torch.float64)
tensor(0.6528, dtype=torch.float64) tensor(0.3629, dtype=torch.float64)
tensor(0.6288, dtype=torch.float64) tensor(0.3552, dtype=torch.float64)
tensor(0.6581, dtype=torch.float64) tensor(0.3473, dtype=torch.float64)
tensor(0.6333, dtype=torch.float64) tensor(0.3395, dtype=torch.float64)
tensor(0.6086, dtype=torch.float64) tensor(0.3319, dtype=torch.float64)
tensor(0.5213, dtype=torch.float64) tensor(0.3247, dtype=torch.float64)
tensor(0.5545, dtype=torch.float64) tensor(0.3174, dtype=torch.float64)
tensor(0.5322, dtype=torch.float64) tensor(0.3102, dtype=torch.float64)
tensor(0.4551, dtype=torch.float64) tensor(0.3034, dtype=torch.float64)
tensor(0.4386, dtype=torch.float64) tensor(0.2967, dtype=torch.float64)
tensor(0.4220, dtype=torch.float64) tensor(0.2900, dtype=torch.float64)
tensor(0.4055, dtype=torch.float64) tensor(0.2835, dtype=torch.float64)
tensor(0.4140, dtype=torch.float64) tensor(0.2769, dtype=torch.float64)
tensor(0.3690, dtype=torch.float64) tensor(0.2706, dtype=torch.float64)
tensor(0.3541, dtype=torch.float64) tensor(0.2645, dtype=torch.float64)
tensor(0.3395, dtype=torch.float64) tensor(0.2584, dtype=torch.float64)
tensor(0.3251, dtype=torch.float64) tensor(0.2524, dtype=torch.float64)
tensor(0.3246, dtype=torch.float64) tensor(0.2464, dtype=torch.float64)
tensor(0.3123, dtype=torch.float64) tensor(0.2405, dtype=torch.float64)
tensor(0.2781, dtype=torch.float64) tensor(0.2349, dtype=torch.float64)
tensor(0.2846, dtype=torch.float64) tensor(0.2292, dtype=torch.float64)
tensor(0.2732, dtype=torch.float64) tensor(0.2237, dtype=torch.float64)
tensor(0.2619, dtype=torch.float64) tensor(0.2183, dtype=torch.float64)
tensor(0.2507, dtype=torch.float64) tensor(0.2130, dtype=torch.float64)
tensor(0.2145, dtype=torch.float64) tensor(0.2079, dtype=torch.float64)
tensor(0.2265, dtype=torch.float64) tensor(0.2028, dtype=torch.float64)
tensor(0.2166, dtype=torch.float64) tensor(0.1979, dtype=torch.float64)
tensor(0.1852, dtype=torch.float64) tensor(0.1931, dtype=torch.float64)
tensor(0.1779, dtype=torch.float64) tensor(0.1885, dtype=torch.float64)
tensor(0.1706, dtype=torch.float64) tensor(0.1839, dtype=torch.float64)
tensor(0.1743, dtype=torch.float64) tensor(0.1793, dtype=torch.float64)
tensor(0.1669, dtype=torch.float64) tensor(0.1748, dtype=torch.float64)
tensor(0.1460, dtype=torch.float64) tensor(0.1705, dtype=torch.float64)
tensor(0.1399, dtype=torch.float64) tensor(0.1663, dtype=torch.float64)
tensor(0.1422, dtype=torch.float64) tensor(0.1621, dtype=torch.float64)
tensor(0.1263, dtype=torch.float64) tensor(0.1581, dtype=torch.float64)
tensor(0.1284, dtype=torch.float64) tensor(0.1540, dtype=torch.float64)
tensor(0.1139, dtype=torch.float64) tensor(0.1501, dtype=torch.float64)
tensor(0.1089, dtype=torch.float64) tensor(0.1464, dtype=torch.float64)
tensor(0.1040, dtype=torch.float64) tensor(0.1427, dtype=torch.float64)
tensor(0.1035, dtype=torch.float64) tensor(0.1389, dtype=torch.float64)
tensor(0.0933, dtype=torch.float64) tensor(0.1354, dtype=torch.float64)
tensor(0.0935, dtype=torch.float64) tensor(0.1318, dtype=torch.float64)
tensor(0.0837, dtype=torch.float64) tensor(0.1284, dtype=torch.float64)
tensor(0.0844, dtype=torch.float64) tensor(0.1250, dtype=torch.float64)
tensor(0.0751, dtype=torch.float64) tensor(0.1218, dtype=torch.float64)
tensor(0.0761, dtype=torch.float64) tensor(0.1185, dtype=torch.float64)
tensor(0.0674, dtype=torch.float64) tensor(0.1154, dtype=torch.float64)
tensor(0.0642, dtype=torch.float64) tensor(0.1124, dtype=torch.float64)
tensor(0.0612, dtype=torch.float64) tensor(0.1094, dtype=torch.float64)
tensor(0.0609, dtype=torch.float64) tensor(0.1064, dtype=torch.float64)
tensor(0.0546, dtype=torch.float64) tensor(0.1036, dtype=torch.float64)
tensor(0.0520, dtype=torch.float64) tensor(0.1008, dtype=torch.float64)
tensor(0.0518, dtype=torch.float64) tensor(0.0980, dtype=torch.float64)
tensor(0.0463, dtype=torch.float64) tensor(0.0954, dtype=torch.float64)
tensor(0.0466, dtype=torch.float64) tensor(0.0928, dtype=torch.float64)
tensor(0.0445, dtype=torch.float64) tensor(0.0902, dtype=torch.float64)
tensor(0.0387, dtype=torch.float64) tensor(0.0877, dtype=torch.float64)
tensor(0.0399, dtype=torch.float64) tensor(0.0853, dtype=torch.float64)
tensor(0.0380, dtype=torch.float64) tensor(0.0829, dtype=torch.float64)
tensor(0.0361, dtype=torch.float64) tensor(0.0806, dtype=torch.float64)
tensor(0.0306, dtype=torch.float64) tensor(0.0784, dtype=torch.float64)
tensor(0.0291, dtype=torch.float64) tensor(0.0762, dtype=torch.float64)
tensor(0.0278, dtype=torch.float64) tensor(0.0741, dtype=torch.float64)
tensor(0.0264, dtype=torch.float64) tensor(0.0720, dtype=torch.float64)
tensor(0.0265, dtype=torch.float64) tensor(0.0700, dtype=torch.float64)
tensor(0.0253, dtype=torch.float64) tensor(0.0680, dtype=torch.float64)
tensor(0.0240, dtype=torch.float64) tensor(0.0660, dtype=torch.float64)
tensor(0.0206, dtype=torch.float64) tensor(0.0641, dtype=torch.float64)
tensor(0.0196, dtype=torch.float64) tensor(0.0623, dtype=torch.float64)
tensor(0.0200, dtype=torch.float64) tensor(0.0605, dtype=torch.float64)
tensor(0.0175, dtype=torch.float64) tensor(0.0588, dtype=torch.float64)
tensor(0.0166, dtype=torch.float64) tensor(0.0571, dtype=torch.float64)
tensor(0.0158, dtype=torch.float64) tensor(0.0555, dtype=torch.float64)
tensor(0.0149, dtype=torch.float64) tensor(0.0539, dtype=torch.float64)
tensor(0.0141, dtype=torch.float64) tensor(0.0523, dtype=torch.float64)
tensor(0.0134, dtype=torch.float64) tensor(0.0508, dtype=torch.float64)
tensor(0.0126, dtype=torch.float64) tensor(0.0493, dtype=torch.float64)
tensor(0.0119, dtype=torch.float64) tensor(0.0478, dtype=torch.float64)
tensor(0.0114, dtype=torch.float64) tensor(0.0464, dtype=torch.float64)
tensor(0.0104, dtype=torch.float64) tensor(0.0450, dtype=torch.float64)
tensor(0.0102, dtype=torch.float64) tensor(0.0436, dtype=torch.float64)
tensor(0.0098, dtype=torch.float64) tensor(0.0423, dtype=torch.float64)
tensor(0.0093, dtype=torch.float64) tensor(0.0410, dtype=torch.float64)
tensor(0.0088, dtype=torch.float64) tensor(0.0397, dtype=torch.float64)
tensor(0.0074, dtype=torch.float64) tensor(0.0385, dtype=torch.float64)
tensor(0.0078, dtype=torch.float64) tensor(0.0373, dtype=torch.float64)
tensor(0.0074, dtype=torch.float64) tensor(0.0362, dtype=torch.float64)
tensor(0.0070, dtype=torch.float64) tensor(0.0351, dtype=torch.float64)
tensor(0.0066, dtype=torch.float64) tensor(0.0340, dtype=torch.float64)
tensor(0.0054, dtype=torch.float64) tensor(0.0330, dtype=torch.float64)
tensor(0.0051, dtype=torch.float64) tensor(0.0320, dtype=torch.float64)
tensor(0.0048, dtype=torch.float64) tensor(0.0310, dtype=torch.float64)
tensor(0.0046, dtype=torch.float64) tensor(0.0300, dtype=torch.float64)
tensor(0.0046, dtype=torch.float64) tensor(0.0291, dtype=torch.float64)
tensor(0.0040, dtype=torch.float64) tensor(0.0282, dtype=torch.float64)
tensor(0.0041, dtype=torch.float64) tensor(0.0273, dtype=torch.float64)
tensor(0.0036, dtype=torch.float64) tensor(0.0264, dtype=torch.float64)
tensor(0.0034, dtype=torch.float64) tensor(0.0256, dtype=torch.float64)
tensor(0.0032, dtype=torch.float64) tensor(0.0248, dtype=torch.float64)
tensor(0.0031, dtype=torch.float64) tensor(0.0240, dtype=torch.float64)
tensor(0.0029, dtype=torch.float64) tensor(0.0232, dtype=torch.float64)
tensor(0.0028, dtype=torch.float64) tensor(0.0225, dtype=torch.float64)
tensor(0.0024, dtype=torch.float64) tensor(0.0218, dtype=torch.float64)
tensor(0.0023, dtype=torch.float64) tensor(0.0211, dtype=torch.float64)
tensor(0.0023, dtype=torch.float64) tensor(0.0204, dtype=torch.float64)
tensor(0.0021, dtype=torch.float64) tensor(0.0197, dtype=torch.float64)
tensor(0.0018, dtype=torch.float64) tensor(0.0191, dtype=torch.float64)
tensor(0.0019, dtype=torch.float64) tensor(0.0184, dtype=torch.float64)
tensor(0.0018, dtype=torch.float64) tensor(0.0178, dtype=torch.float64)
tensor(0.0017, dtype=torch.float64) tensor(0.0172, dtype=torch.float64)
tensor(0.0014, dtype=torch.float64) tensor(0.0167, dtype=torch.float64)
tensor(0.0014, dtype=torch.float64) tensor(0.0161, dtype=torch.float64)
tensor(0.0014, dtype=torch.float64) tensor(0.0156, dtype=torch.float64)
tensor(0.0011, dtype=torch.float64) tensor(0.0151, dtype=torch.float64)
tensor(0.0012, dtype=torch.float64) tensor(0.0146, dtype=torch.float64)
tensor(0.0010, dtype=torch.float64) tensor(0.0141, dtype=torch.float64)
tensor(0.0010, dtype=torch.float64) tensor(0.0136, dtype=torch.float64)
tensor(0.0009, dtype=torch.float64) tensor(0.0132, dtype=torch.float64)
tensor(0.0008, dtype=torch.float64) tensor(0.0127, dtype=torch.float64)
tensor(0.0008, dtype=torch.float64) tensor(0.0123, dtype=torch.float64)
tensor(0.0008, dtype=torch.float64) tensor(0.0119, dtype=torch.float64)
tensor(0.0007, dtype=torch.float64) tensor(0.0114, dtype=torch.float64)
tensor(0.0007, dtype=torch.float64) tensor(0.0110, dtype=torch.float64)
tensor(0.0006, dtype=torch.float64) tensor(0.0107, dtype=torch.float64)
tensor(0.0006, dtype=torch.float64) tensor(0.0103, dtype=torch.float64)
tensor(0.0006, dtype=torch.float64) tensor(0.0099, dtype=torch.float64)

Example-25: Normalize

[1]:
# In this example normalized objective construction is illustrated
[2]:
# Import

import torch
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader
torch.set_printoptions(linewidth=128)

import matplotlib
from matplotlib import pyplot as plt
matplotlib.rcParams['text.usetex'] = True

from twiss import twiss

from ndmap.signature import chop
from ndmap.evaluate import evaluate
from ndmap.pfp import parametric_fixed_point

from model.library.drift import Drift
from model.library.quadrupole import Quadrupole
from model.library.sextupole import Sextupole
from model.library.dipole import Dipole
from model.library.line import Line

from model.command.wrapper import group
from model.command.wrapper import forward
from model.command.wrapper import inverse
from model.command.wrapper import normalize
from model.command.wrapper import Wrapper
[3]:
# Define simple FODO based lattice using nested lines

DR = Drift('DR', 0.25)
BM = Dipole('BM', 3.50, torch.pi/4.0)

QF_A = Quadrupole('QF_A', 0.5, +0.20)
QD_A = Quadrupole('QD_A', 0.5, -0.19)
QF_B = Quadrupole('QF_B', 0.5, +0.20)
QD_B = Quadrupole('QD_B', 0.5, -0.19)
QF_C = Quadrupole('QF_C', 0.5, +0.20)
QD_C = Quadrupole('QD_C', 0.5, -0.19)
QF_D = Quadrupole('QF_D', 0.5, +0.20)
QD_D = Quadrupole('QD_D', 0.5, -0.19)

SF_A = Sextupole('SF_A', 0.25, 0.00)
SD_A = Sextupole('SD_A', 0.25, 0.00)
SF_B = Sextupole('SF_B', 0.25, 0.00)
SD_B = Sextupole('SD_B', 0.25, 0.00)
SF_C = Sextupole('SF_C', 0.25, 0.00)
SD_C = Sextupole('SD_C', 0.25, 0.00)
SF_D = Sextupole('SF_D', 0.25, 0.00)
SD_D = Sextupole('SD_D', 0.25, 0.00)

FODO_A = Line('FODO_A', [QF_A, DR, SF_A, DR, BM, DR, SD_A, DR, QD_A, QD_A, DR, SD_A, DR, BM, DR, SF_A, DR, QF_A], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_B = Line('FODO_B', [QF_B, DR, SF_B, DR, BM, DR, SD_B, DR, QD_B, QD_B, DR, SD_B, DR, BM, DR, SF_B, DR, QF_B], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_C = Line('FODO_C', [QF_C, DR, SF_C, DR, BM, DR, SD_C, DR, QD_C, QD_C, DR, SD_C, DR, BM, DR, SF_C, DR, QF_C], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_D = Line('FODO_D', [QF_D, DR, SF_D, DR, BM, DR, SD_D, DR, QD_D, QD_D, DR, SD_D, DR, BM, DR, SF_D, DR, QF_D], propagate=True, dp=0.0, exact=False, output=False, matrix=False)

RING = Line('RING', [FODO_A, FODO_B, FODO_C, FODO_D], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
[4]:
# Set parametric mapping

ring, *_ = group(RING, 'FODO_A', 'FODO_D', ('ms', ['Sextupole'], None, None), ('dp', None, None, None), root=True)
[5]:
# Construct normalized function

fn = normalize(ring, [(None, None), (-10.0, 10.0), (-0.01, 0.01)])

# Compare with original

fp = torch.tensor([0.001, 0.0005, -0.010, 0.0025], dtype=torch.float64)
ms = torch.tensor([1.0, -1.0, 0.5, 2.0, 4.0, -5.0, -1.0, 3.0], dtype=torch.float64)
dp = torch.tensor([0.005], dtype=torch.float64)

print(ring(fp, ms, dp))
print(fn(*forward([fp, ms, dp],  [(None, None), (-10.0, 10.0), (-0.01, 0.01)])))
tensor([ 0.0157, -0.0006, -0.0189, -0.0032], dtype=torch.float64)
tensor([ 0.0157, -0.0006, -0.0189, -0.0032], dtype=torch.float64)
[6]:
# Set deviation parameters

fp = torch.tensor(4*[0.0], dtype=torch.float64)
ms = torch.tensor(8*[0.0], dtype=torch.float64)
dp = torch.tensor([0.0], dtype=torch.float64)
[7]:
# Define parametric chomaticity function

# Compute parametric fixed point (first order dispersion)

pfp, *_ = parametric_fixed_point((0, 1), fp, [ms, dp], ring)
chop(pfp)

# Define ring around parametric fixed point

def mapping(state, ms, dp):
    return ring(state + evaluate(pfp, [ms, dp]), ms, dp) - evaluate(pfp, [ms, dp])

# Define tunes

def tune(ms, dp):
    matrix = torch.func.jacrev(mapping)(fp, ms, dp)
    tunes, *_ = twiss(matrix)
    return tunes

# Define chromaticity

def chromaticity(ms):
    return torch.func.jacrev(tune, 1)(ms, dp).squeeze()

# Compute natural chromaticity

print(chromaticity(ms))
tensor([-2.0649, -0.8260], dtype=torch.float64)
[8]:
# Chromaticity can be corrected in a single step

# Compute starting values

psix, psiy = chromaticity(ms)

# Set target values

psix_target = torch.tensor(5.0, dtype=torch.float64)
psiy_target = torch.tensor(5.0, dtype=torch.float64)

# Perform correction

dpsix = psix - psix_target
dpsiy = psiy - psiy_target

solution = - torch.linalg.pinv((torch.func.jacrev(chromaticity)(ms)).squeeze()) @ torch.stack([dpsix, dpsiy])
print(solution)

# Test solution

print(chromaticity(solution))
tensor([ 0.7439, -1.2084,  0.7439, -1.2084,  0.7439, -1.2084,  0.7439, -1.2084], dtype=torch.float64)
tensor([5.0000, 5.0000], dtype=torch.float64)
[9]:
# Optimization (wrapping objective funtion and normalization)

# Set model parameters
# Parameters are not cloned inside the module on initialization, values will change during optimization!

ms = torch.tensor(8*[0.0], dtype=torch.float64)
ms, *_ = forward([ms], [(-10, 10)])

# Define scalar objective function

def objective(ms):
    psix, psiy = chromaticity(ms)
    return ((psix - psix_target)**2 + (psiy - psiy_target)**2).sqrt()

print(objective(solution))

# Define normalized objective

objective = normalize(objective, [(-10.0, 10.0)])

print(objective(*forward([solution], [(-10, 10)])))


# Set model (forward returns evaluated objective)

model = Wrapper(objective, ms)

# Set optimizer

optimizer = torch.optim.Adam(model.parameters(), lr=1.0E-3)

# Perfom optimization

epochs = 128
for epoch in range(epochs):

    # Evaluate model
    error = model()

    # Compute derivatives
    error.backward()

    # Perform optimization step
    optimizer.step()

    # Set gradient to zero
    optimizer.zero_grad()

    # Verbose
    knobs, *_ = [*model.parameters()]
    knobs, *_ = inverse([knobs], [(-10, 10)])
    print(error.detach(), (knobs.detach() - solution).norm())
tensor(5.6871e-15, dtype=torch.float64)
tensor(1.1580e-14, dtype=torch.float64)
tensor(9.1573, dtype=torch.float64) tensor(2.7830, dtype=torch.float64)
tensor(8.9651, dtype=torch.float64) tensor(2.7280, dtype=torch.float64)
tensor(8.7737, dtype=torch.float64) tensor(2.6732, dtype=torch.float64)
tensor(8.5832, dtype=torch.float64) tensor(2.6184, dtype=torch.float64)
tensor(8.3937, dtype=torch.float64) tensor(2.5637, dtype=torch.float64)
tensor(8.2052, dtype=torch.float64) tensor(2.5090, dtype=torch.float64)
tensor(8.0177, dtype=torch.float64) tensor(2.4545, dtype=torch.float64)
tensor(7.8314, dtype=torch.float64) tensor(2.4000, dtype=torch.float64)
tensor(7.6464, dtype=torch.float64) tensor(2.3456, dtype=torch.float64)
tensor(7.4627, dtype=torch.float64) tensor(2.2914, dtype=torch.float64)
tensor(7.2804, dtype=torch.float64) tensor(2.2373, dtype=torch.float64)
tensor(7.0995, dtype=torch.float64) tensor(2.1833, dtype=torch.float64)
tensor(6.9202, dtype=torch.float64) tensor(2.1294, dtype=torch.float64)
tensor(6.7426, dtype=torch.float64) tensor(2.0758, dtype=torch.float64)
tensor(6.5666, dtype=torch.float64) tensor(2.0222, dtype=torch.float64)
tensor(6.3924, dtype=torch.float64) tensor(1.9689, dtype=torch.float64)
tensor(6.2200, dtype=torch.float64) tensor(1.9158, dtype=torch.float64)
tensor(6.0495, dtype=torch.float64) tensor(1.8628, dtype=torch.float64)
tensor(5.8808, dtype=torch.float64) tensor(1.8101, dtype=torch.float64)
tensor(5.7141, dtype=torch.float64) tensor(1.7577, dtype=torch.float64)
tensor(5.5492, dtype=torch.float64) tensor(1.7055, dtype=torch.float64)
tensor(5.3862, dtype=torch.float64) tensor(1.6536, dtype=torch.float64)
tensor(5.2250, dtype=torch.float64) tensor(1.6019, dtype=torch.float64)
tensor(5.0655, dtype=torch.float64) tensor(1.5506, dtype=torch.float64)
tensor(4.9077, dtype=torch.float64) tensor(1.4996, dtype=torch.float64)
tensor(4.7514, dtype=torch.float64) tensor(1.4489, dtype=torch.float64)
tensor(4.5965, dtype=torch.float64) tensor(1.3986, dtype=torch.float64)
tensor(4.4428, dtype=torch.float64) tensor(1.3486, dtype=torch.float64)
tensor(4.2902, dtype=torch.float64) tensor(1.2990, dtype=torch.float64)
tensor(4.1384, dtype=torch.float64) tensor(1.2498, dtype=torch.float64)
tensor(3.9873, dtype=torch.float64) tensor(1.2009, dtype=torch.float64)
tensor(3.8366, dtype=torch.float64) tensor(1.1523, dtype=torch.float64)
tensor(3.6862, dtype=torch.float64) tensor(1.1041, dtype=torch.float64)
tensor(3.5359, dtype=torch.float64) tensor(1.0563, dtype=torch.float64)
tensor(3.3857, dtype=torch.float64) tensor(1.0088, dtype=torch.float64)
tensor(3.2353, dtype=torch.float64) tensor(0.9616, dtype=torch.float64)
tensor(3.0848, dtype=torch.float64) tensor(0.9148, dtype=torch.float64)
tensor(2.9343, dtype=torch.float64) tensor(0.8682, dtype=torch.float64)
tensor(2.7836, dtype=torch.float64) tensor(0.8220, dtype=torch.float64)
tensor(2.6331, dtype=torch.float64) tensor(0.7759, dtype=torch.float64)
tensor(2.4828, dtype=torch.float64) tensor(0.7301, dtype=torch.float64)
tensor(2.3329, dtype=torch.float64) tensor(0.6845, dtype=torch.float64)
tensor(2.1835, dtype=torch.float64) tensor(0.6390, dtype=torch.float64)
tensor(2.0349, dtype=torch.float64) tensor(0.5936, dtype=torch.float64)
tensor(1.8870, dtype=torch.float64) tensor(0.5482, dtype=torch.float64)
tensor(1.7399, dtype=torch.float64) tensor(0.5028, dtype=torch.float64)
tensor(1.5934, dtype=torch.float64) tensor(0.4573, dtype=torch.float64)
tensor(1.4474, dtype=torch.float64) tensor(0.4116, dtype=torch.float64)
tensor(1.3016, dtype=torch.float64) tensor(0.3657, dtype=torch.float64)
tensor(1.1557, dtype=torch.float64) tensor(0.3195, dtype=torch.float64)
tensor(1.0093, dtype=torch.float64) tensor(0.2729, dtype=torch.float64)
tensor(0.8622, dtype=torch.float64) tensor(0.2261, dtype=torch.float64)
tensor(0.7146, dtype=torch.float64) tensor(0.1791, dtype=torch.float64)
tensor(0.5667, dtype=torch.float64) tensor(0.1320, dtype=torch.float64)
tensor(0.4196, dtype=torch.float64) tensor(0.0852, dtype=torch.float64)
tensor(0.2748, dtype=torch.float64) tensor(0.0393, dtype=torch.float64)
tensor(0.1340, dtype=torch.float64) tensor(0.0068, dtype=torch.float64)
tensor(0.0262, dtype=torch.float64) tensor(0.0432, dtype=torch.float64)
tensor(0.1620, dtype=torch.float64) tensor(0.0727, dtype=torch.float64)
tensor(0.2647, dtype=torch.float64) tensor(0.0948, dtype=torch.float64)
tensor(0.3263, dtype=torch.float64) tensor(0.1111, dtype=torch.float64)
tensor(0.3638, dtype=torch.float64) tensor(0.1224, dtype=torch.float64)
tensor(0.3894, dtype=torch.float64) tensor(0.1289, dtype=torch.float64)
tensor(0.4067, dtype=torch.float64) tensor(0.1301, dtype=torch.float64)
tensor(0.4121, dtype=torch.float64) tensor(0.1261, dtype=torch.float64)
tensor(0.4008, dtype=torch.float64) tensor(0.1169, dtype=torch.float64)
tensor(0.3711, dtype=torch.float64) tensor(0.1032, dtype=torch.float64)
tensor(0.3262, dtype=torch.float64) tensor(0.0861, dtype=torch.float64)
tensor(0.2728, dtype=torch.float64) tensor(0.0667, dtype=torch.float64)
tensor(0.2177, dtype=torch.float64) tensor(0.0458, dtype=torch.float64)
tensor(0.1600, dtype=torch.float64) tensor(0.0233, dtype=torch.float64)
tensor(0.0867, dtype=torch.float64) tensor(0.0035, dtype=torch.float64)
tensor(0.0227, dtype=torch.float64) tensor(0.0224, dtype=torch.float64)
tensor(0.0708, dtype=torch.float64) tensor(0.0378, dtype=torch.float64)
tensor(0.1198, dtype=torch.float64) tensor(0.0465, dtype=torch.float64)
tensor(0.1470, dtype=torch.float64) tensor(0.0496, dtype=torch.float64)
tensor(0.1569, dtype=torch.float64) tensor(0.0479, dtype=torch.float64)
tensor(0.1535, dtype=torch.float64) tensor(0.0422, dtype=torch.float64)
tensor(0.1370, dtype=torch.float64) tensor(0.0329, dtype=torch.float64)
tensor(0.1064, dtype=torch.float64) tensor(0.0205, dtype=torch.float64)
tensor(0.0647, dtype=torch.float64) tensor(0.0052, dtype=torch.float64)
tensor(0.0201, dtype=torch.float64) tensor(0.0155, dtype=torch.float64)
tensor(0.0578, dtype=torch.float64) tensor(0.0268, dtype=torch.float64)
tensor(0.0859, dtype=torch.float64) tensor(0.0327, dtype=torch.float64)
tensor(0.1052, dtype=torch.float64) tensor(0.0344, dtype=torch.float64)
tensor(0.1163, dtype=torch.float64) tensor(0.0319, dtype=torch.float64)
tensor(0.1068, dtype=torch.float64) tensor(0.0257, dtype=torch.float64)
tensor(0.0814, dtype=torch.float64) tensor(0.0165, dtype=torch.float64)
tensor(0.0557, dtype=torch.float64) tensor(0.0028, dtype=torch.float64)
tensor(0.0120, dtype=torch.float64) tensor(0.0180, dtype=torch.float64)
tensor(0.0825, dtype=torch.float64) tensor(0.0273, dtype=torch.float64)
tensor(0.1034, dtype=torch.float64) tensor(0.0305, dtype=torch.float64)
tensor(0.0967, dtype=torch.float64) tensor(0.0315, dtype=torch.float64)
tensor(0.1101, dtype=torch.float64) tensor(0.0296, dtype=torch.float64)
tensor(0.1152, dtype=torch.float64) tensor(0.0230, dtype=torch.float64)
tensor(0.0848, dtype=torch.float64) tensor(0.0143, dtype=torch.float64)
tensor(0.0472, dtype=torch.float64) tensor(0.0074, dtype=torch.float64)
tensor(0.0454, dtype=torch.float64) tensor(0.0103, dtype=torch.float64)
tensor(0.0325, dtype=torch.float64) tensor(0.0179, dtype=torch.float64)
tensor(0.0641, dtype=torch.float64) tensor(0.0188, dtype=torch.float64)
tensor(0.0619, dtype=torch.float64) tensor(0.0155, dtype=torch.float64)
tensor(0.0520, dtype=torch.float64) tensor(0.0098, dtype=torch.float64)
tensor(0.0419, dtype=torch.float64) tensor(0.0017, dtype=torch.float64)
tensor(0.0103, dtype=torch.float64) tensor(0.0091, dtype=torch.float64)
tensor(0.0345, dtype=torch.float64) tensor(0.0123, dtype=torch.float64)
tensor(0.0390, dtype=torch.float64) tensor(0.0117, dtype=torch.float64)
tensor(0.0417, dtype=torch.float64) tensor(0.0071, dtype=torch.float64)
tensor(0.0244, dtype=torch.float64) tensor(0.0044, dtype=torch.float64)
tensor(0.0279, dtype=torch.float64) tensor(0.0079, dtype=torch.float64)
tensor(0.0253, dtype=torch.float64) tensor(0.0103, dtype=torch.float64)
tensor(0.0361, dtype=torch.float64) tensor(0.0074, dtype=torch.float64)
tensor(0.0241, dtype=torch.float64) tensor(0.0023, dtype=torch.float64)
tensor(0.0138, dtype=torch.float64) tensor(0.0086, dtype=torch.float64)
tensor(0.0475, dtype=torch.float64) tensor(0.0110, dtype=torch.float64)
tensor(0.0466, dtype=torch.float64) tensor(0.0118, dtype=torch.float64)
tensor(0.0417, dtype=torch.float64) tensor(0.0105, dtype=torch.float64)
tensor(0.0462, dtype=torch.float64) tensor(0.0033, dtype=torch.float64)
tensor(0.0105, dtype=torch.float64) tensor(0.0076, dtype=torch.float64)
tensor(0.0357, dtype=torch.float64) tensor(0.0115, dtype=torch.float64)
tensor(0.0363, dtype=torch.float64) tensor(0.0128, dtype=torch.float64)
tensor(0.0477, dtype=torch.float64) tensor(0.0097, dtype=torch.float64)
tensor(0.0349, dtype=torch.float64) tensor(0.0050, dtype=torch.float64)
tensor(0.0254, dtype=torch.float64) tensor(0.0042, dtype=torch.float64)
tensor(0.0137, dtype=torch.float64) tensor(0.0082, dtype=torch.float64)
tensor(0.0353, dtype=torch.float64) tensor(0.0063, dtype=torch.float64)
tensor(0.0205, dtype=torch.float64) tensor(0.0052, dtype=torch.float64)
tensor(0.0327, dtype=torch.float64) tensor(0.0035, dtype=torch.float64)
tensor(0.0119, dtype=torch.float64) tensor(0.0082, dtype=torch.float64)
[10]:
# Compare

print(solution)
print(*inverse([ms], [(-10, 10)]))
tensor([ 0.7439, -1.2084,  0.7439, -1.2084,  0.7439, -1.2084,  0.7439, -1.2084], dtype=torch.float64)
tensor([ 0.7412, -1.2115,  0.7412, -1.2115,  0.7412, -1.2115,  0.7412, -1.2115], dtype=torch.float64)

Example-26: Transformation

[1]:
# In this example another wrappers are used to construct parametric transformations between elements
# Given two element a transformation can be constructed from the first element enterence frame to the second element exit frame
# If the first element appears after the second one in the line, inverse transformation is constructed
# Note, these transformations are given around initial reference orbit
[2]:
# Import

import torch

import matplotlib
from matplotlib import pyplot as plt
matplotlib.rcParams['text.usetex'] = True

from twiss import twiss

from ndmap.signature import chop
from ndmap.evaluate import evaluate
from ndmap.pfp import parametric_fixed_point

from model.library.drift import Drift
from model.library.quadrupole import Quadrupole
from model.library.sextupole import Sextupole
from model.library.dipole import Dipole
from model.library.line import Line

from model.command.wrapper import group
[3]:
# Define simple FODO based lattice using nested lines
# Note, all elements have unique names

QF_A = Quadrupole('QF_A', 1.0, +0.20)
QD_A = Quadrupole('QD_A', 1.0, -0.19)
QF_B = Quadrupole('QF_B', 1.0, +0.20)
QD_B = Quadrupole('QD_B', 1.0, -0.19)
QF_C = Quadrupole('QF_C', 1.0, +0.20)
QD_C = Quadrupole('QD_C', 1.0, -0.19)
QF_D = Quadrupole('QF_D', 1.0, +0.20)
QD_D = Quadrupole('QD_D', 1.0, -0.19)

SF1_A = Sextupole('SF1_A', 0.25, 0.00)
SD1_A = Sextupole('SD1_A', 0.25, 0.00)
SF2_A = Sextupole('SF2_A', 0.25, 0.00)
SD2_A = Sextupole('SD2_A', 0.25, 0.00)
SF1_B = Sextupole('SF1_B', 0.25, 0.00)
SD1_B = Sextupole('SD1_B', 0.25, 0.00)
SF2_B = Sextupole('SF2_B', 0.25, 0.00)
SD2_B = Sextupole('SD2_B', 0.25, 0.00)
SF1_C = Sextupole('SF1_C', 0.25, 0.00)
SD1_C = Sextupole('SD1_C', 0.25, 0.00)
SF2_C = Sextupole('SF2_C', 0.25, 0.00)
SD2_C = Sextupole('SD2_C', 0.25, 0.00)
SF1_D = Sextupole('SF1_D', 0.25, 0.00)
SD1_D = Sextupole('SD1_D', 0.25, 0.00)
SF2_D = Sextupole('SF2_D', 0.25, 0.00)
SD2_D = Sextupole('SD2_D', 0.25, 0.00)

BM1_A = Dipole('BM1_A', 3.50, torch.pi/4.0)
BM2_A = Dipole('BM2_A', 3.50, torch.pi/4.0)
BM1_B = Dipole('BM1_B', 3.50, torch.pi/4.0)
BM2_B = Dipole('BM2_B', 3.50, torch.pi/4.0)
BM1_C = Dipole('BM1_C', 3.50, torch.pi/4.0)
BM2_C = Dipole('BM2_C', 3.50, torch.pi/4.0)
BM1_D = Dipole('BM1_D', 3.50, torch.pi/4.0)
BM2_D = Dipole('BM2_D', 3.50, torch.pi/4.0)

DR1_A = Drift('DR1_A', 0.25)
DR2_A = Drift('DR2_A', 0.25)
DR3_A = Drift('DR3_A', 0.25)
DR4_A = Drift('DR4_A', 0.25)
DR5_A = Drift('DR5_A', 0.25)
DR6_A = Drift('DR6_A', 0.25)
DR7_A = Drift('DR7_A', 0.25)
DR1_B = Drift('DR1_B', 0.25)
DR2_B = Drift('DR2_B', 0.25)
DR3_B = Drift('DR3_B', 0.25)
DR4_B = Drift('DR4_B', 0.25)
DR5_B = Drift('DR5_B', 0.25)
DR6_B = Drift('DR6_B', 0.25)
DR7_B = Drift('DR7_B', 0.25)
DR1_C = Drift('DR1_C', 0.25)
DR2_C = Drift('DR2_C', 0.25)
DR3_C = Drift('DR3_C', 0.25)
DR4_C = Drift('DR4_C', 0.25)
DR5_C = Drift('DR5_C', 0.25)
DR6_C = Drift('DR6_C', 0.25)
DR7_C = Drift('DR7_C', 0.25)
DR1_D = Drift('DR1_D', 0.25)
DR2_D = Drift('DR2_D', 0.25)
DR3_D = Drift('DR3_D', 0.25)
DR4_D = Drift('DR4_D', 0.25)
DR5_D = Drift('DR5_D', 0.25)
DR6_D = Drift('DR6_D', 0.25)
DR7_D = Drift('DR7_D', 0.25)

FODO_A = Line('FODO_A', [QF_A, DR1_A, SF1_A, DR2_A, BM1_A, DR3_A, SD1_A, DR3_A, QD_A, DR4_A, SD2_A, DR5_A, BM2_A, DR6_A, SF2_A, DR7_A], propagate=True)
FODO_B = Line('FODO_B', [QF_B, DR1_B, SF1_B, DR2_B, BM1_B, DR3_B, SD1_B, DR3_B, QD_B, DR4_B, SD2_B, DR5_B, BM2_B, DR6_B, SF2_B, DR7_B], propagate=True)
FODO_C = Line('FODO_C', [QF_C, DR1_C, SF1_C, DR2_C, BM1_C, DR3_C, SD1_C, DR3_C, QD_C, DR4_C, SD2_C, DR5_C, BM2_C, DR6_C, SF2_C, DR7_C], propagate=True)
FODO_D = Line('FODO_D', [QF_D, DR1_D, SF1_D, DR2_D, BM1_D, DR3_D, SD1_D, DR3_D, QD_D, DR4_D, SD2_D, DR5_D, BM2_D, DR6_D, SF2_D, DR7_D], propagate=True)

RING = Line('RING', [FODO_A, FODO_B, FODO_C, FODO_D], propagate=True, dp=0.0, exact=False, output=False, matrix=False)

RING.flatten()
[4]:
# Create parametric transformation from one element to another and its inverse

probe = 'SD2_A'
other = 'SF1_D'

forward, *_ = group(RING, probe, other, ('kn', ['Quadrupole'], None, None), ('ms', ['Sextupole'], None, None), ('dp', None, None, None), root=True, alignment=False)
inverse, *_ = group(RING, other, probe, ('kn', ['Quadrupole'], None, None), ('ms', ['Sextupole'], None, None), ('dp', None, None, None), root=True, alignment=False)
[5]:
# Test propagation and inverse transformation

state = torch.tensor([0.001, 0.005, -0.005, 0.001], dtype=torch.float64)

kn = 1.0E-3*torch.randn( 8, dtype=torch.float64)
ms = 1.0E-3*torch.randn(16, dtype=torch.float64)
dp = torch.tensor([0.001], dtype=torch.float64)

print(local := state.clone())
print(local := forward(local, kn, ms, dp))
print(local := inverse(local, kn, ms, dp))
print(state - local)

tensor([ 0.0010,  0.0050, -0.0050,  0.0010], dtype=torch.float64)
tensor([-0.0041,  0.0022,  0.0047, -0.0001], dtype=torch.float64)
tensor([ 0.0010,  0.0050, -0.0050,  0.0010], dtype=torch.float64)
tensor([ 4.3368e-19, -8.6736e-19, -4.3368e-18,  2.1684e-19],
       dtype=torch.float64)
[6]:
# Test derivatives

state = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64)

kn = torch.zeros( 8, dtype=torch.float64)
ms = torch.zeros(16, dtype=torch.float64)
dp = torch.tensor([0.0], dtype=torch.float64)

# Transport matrix

print(torch.func.jacrev(forward)(state, kn, ms, dp).inverse())
print(torch.func.jacrev(inverse)(state, kn, ms, dp))
print()

# Derivatives of transport matrix trace with respect to quadrupole deviations

def matrix(kn, ms, dp):
    return torch.func.jacrev(forward)(state, kn, ms, dp).trace()

print(torch.func.jacrev(matrix)(kn, ms, dp))
print()
tensor([[ 0.3917,  0.5359,  0.0000,  0.0000],
        [ 0.4481,  3.1656,  0.0000,  0.0000],
        [-0.0000, -0.0000, -1.1626, -3.2206],
        [ 0.0000,  0.0000,  0.2044, -0.2939]], dtype=torch.float64)
tensor([[ 0.3917,  0.5359,  0.0000,  0.0000],
        [ 0.4481,  3.1656,  0.0000,  0.0000],
        [ 0.0000,  0.0000, -1.1626, -3.2206],
        [ 0.0000,  0.0000,  0.2044, -0.2939]], dtype=torch.float64)

tensor([ 0.0000,  0.0000,  2.6498, 39.4992, 30.9318, 39.5209,  2.0561,  0.0000],
       dtype=torch.float64)

Example-27: Orbit (fixed point computation)

[1]:
# In this example computation of fixed points is illustrated
# Fixed points are computed for given initial guess using Newton root search method
# Closed orbit is computed
# Special case of period one stable (elliptic) fixed point corresponding to center manifold
# Also, period five fixed point is computed (restricted to horizontal plane)
[2]:
# Import

import torch
torch.set_printoptions(linewidth=128)

import matplotlib
from matplotlib import pyplot as plt
matplotlib.rcParams['text.usetex'] = True

from twiss import twiss

from ndmap.signature import chop
from ndmap.evaluate import evaluate
from ndmap.pfp import parametric_fixed_point

from ndmap.pfp import clean_point
from ndmap.pfp import chain_point
from ndmap.pfp import matrix

from model.library.drift import Drift
from model.library.quadrupole import Quadrupole
from model.library.sextupole import Sextupole
from model.library.dipole import Dipole
from model.library.line import Line

from model.command.wrapper import group
from model.command.orbit import orbit
[3]:
# Define simple FODO based lattice using nested lines

DR = Drift('DR', 0.25)
BM = Dipole('BM', 3.50, torch.pi/4.0)

QF_A = Quadrupole('QF_A', 0.5, +0.20)
QD_A = Quadrupole('QD_A', 0.5, -0.19)
QF_B = Quadrupole('QF_B', 0.5, +0.20)
QD_B = Quadrupole('QD_B', 0.5, -0.19)
QF_C = Quadrupole('QF_C', 0.5, +0.20)
QD_C = Quadrupole('QD_C', 0.5, -0.19)
QF_D = Quadrupole('QF_D', 0.5, +0.20)
QD_D = Quadrupole('QD_D', 0.5, -0.19)

SF_A = Sextupole('SF_A', 0.25, 0.00)
SD_A = Sextupole('SD_A', 0.25, 0.00)
SF_B = Sextupole('SF_B', 0.25, 0.00)
SD_B = Sextupole('SD_B', 0.25, 0.00)
SF_C = Sextupole('SF_C', 0.25, 0.00)
SD_C = Sextupole('SD_C', 0.25, 0.00)
SF_D = Sextupole('SF_D', 0.25, 0.00)
SD_D = Sextupole('SD_D', 0.25, 0.00)

FODO_A = Line('FODO_A', [QF_A, DR, SF_A, DR, BM, DR, SD_A, DR, QD_A, QD_A, DR, SD_A, DR, BM, DR, SF_A, DR, QF_A], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_B = Line('FODO_B', [QF_B, DR, SF_B, DR, BM, DR, SD_B, DR, QD_B, QD_B, DR, SD_B, DR, BM, DR, SF_B, DR, QF_B], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_C = Line('FODO_C', [QF_C, DR, SF_C, DR, BM, DR, SD_C, DR, QD_C, QD_C, DR, SD_C, DR, BM, DR, SF_C, DR, QF_C], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_D = Line('FODO_D', [QF_D, DR, SF_D, DR, BM, DR, SD_D, DR, QD_D, QD_D, DR, SD_D, DR, BM, DR, SF_D, DR, QF_D], propagate=True, dp=0.0, exact=False, output=False, matrix=False)

RING = Line('RING', [FODO_A, FODO_B, FODO_C, FODO_D], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
[4]:
# Correct chromaticity

# Set parametric mapping

ring, *_ = group(RING, 'FODO_A', 'FODO_D', ('ms', ['Sextupole'], None, None), ('dp', None, None, None), root=True, alignment=False)

# Set deviation parameters

fp = torch.tensor(4*[0.0], dtype=torch.float64)
dp = torch.tensor([0.0], dtype=torch.float64)
ms = torch.tensor(8*[0.0], dtype=torch.float64)


# Compute first order parametric fixed point with respect to momentum deviation

pfp, *_ = parametric_fixed_point((0, 1), fp, [ms, dp], ring)
chop(pfp)

# Define ring around fixed point

def mapping(state, ms, dp):
    return ring(state + evaluate(pfp, [ms, dp]), ms, dp) - evaluate(pfp, [ms, dp])

# Tune

def tune(ms, dp):
    matrix = torch.func.jacrev(mapping)(fp, ms, dp)
    tunes, *_ = twiss(matrix)
    return tunes

# Chromaticity

def chromaticity(ms):
    return torch.func.jacrev(tune, 1)(ms, dp)

# Initial chomaticity values

psix, psiy = chromaticity(ms).squeeze()

# Define target chomaticity values

psix_target = torch.tensor(5.0, dtype=torch.float64)
psiy_target = torch.tensor(5.0, dtype=torch.float64)

# Perform correction

dpsix = psix - psix_target
dpsiy = psiy - psiy_target

# Set solution

solution = - torch.linalg.pinv((torch.func.jacrev(chromaticity)(ms)).squeeze()) @ torch.stack([dpsix, dpsiy])

# Set sextupoles
# Note, ring function in not effected

SF_A.ms, SD_A.ms, SF_B.ms, SD_B.ms, SF_C.ms, SD_C.ms, SF_D.ms, SD_D.ms = solution.tolist()

# Check chromaticity

print(chromaticity(solution).squeeze())

# Plot tunes vs momentum deviation

nux, nuy = tune(solution, dp)

dps = torch.linspace(-5.0E-3, 5.0E-3, 16, dtype=torch.float64)
nuxs, nuys = torch.stack([tune(solution, dp) for dp in dps.reshape(-1, 1)]).T

plt.figure(figsize=(16, 4))
plt.plot(dps.cpu().numpy(), (nux + psix_target*dps).cpu().numpy(), color='red', linestyle='dashed')
plt.scatter(dps.cpu().numpy(), nuxs.cpu().numpy(), color='black', marker='x')
plt.plot(dps.cpu().numpy(), (nuy + psiy_target*dps).cpu().numpy(), color='blue', linestyle='dashed')
plt.scatter(dps.cpu().numpy(), nuys.cpu().numpy(), color='black', marker='x')
plt.tight_layout()
plt.show()
tensor([5.0000, 5.0000], dtype=torch.float64)
../_images/examples_model_385_1.png
[5]:
# Generate and plot phase space trajectories

qx = torch.linspace(0.10, 0.4, 16, dtype=torch.float64)
px = torch.zeros_like(qx)
qy = torch.zeros_like(qx)
py = torch.zeros_like(qx)

state = torch.stack([qx, px, qy, py]).T
table = []

for _ in range(2**10):
    state = torch.vmap(RING)(state)
    table.append(state)

qx, px, *_ = torch.stack(table).swapaxes(0, -1)

plt.figure(figsize=(6, 6))
plt.scatter(qx.cpu().numpy(), px.cpu().numpy(), s=1, color='black')
plt.xlim(-1.0, 0.5)
plt.ylim(-0.075, 0.075)
plt.tight_layout()
plt.show()
../_images/examples_model_386_0.png
[6]:
# Compute closed orbit (period one fixed point)

# Set initial guess

guess = 1.0E-3*torch.tensor([1.0, -1.0, 1.0, -1.0], dtype=torch.float64)

# Compute using line (deviation parameters are passed directly, wrap with lambda lambda state: RING(state, data=data))

point = orbit(RING, guess, limit=8, epsilon=1.0E-6)

print(point)
print()

# Compute using mapping
# Note, need to update to account for change in sextupoles

ring, *_ = group(RING, 'FODO_A', 'FODO_D', ('ms', ['Sextupole'], None, None), ('dp', None, None, None), root=True, alignment=False)
point = orbit(ring, guess, ms, dp, limit=8, epsilon=1.0E-6)

print(point)
print()

# Compute matrix around closed orbit

print(torch.func.jacrev(RING)(point))
print()

print(matrix(1, ring, point, ms, dp, jacobian=torch.func.jacrev))
print()

# Classify fixed point

values, _ = torch.linalg.eig(matrix(1, ring, point, ms, dp, jacobian=torch.func.jacrev))
print(values.log().real)
print()
tensor([ 2.9857e-18, -2.0620e-19, -2.9123e-19, -3.0706e-20], dtype=torch.float64)

tensor([ 2.9857e-18, -2.0620e-19, -2.9122e-19, -3.0706e-20], dtype=torch.float64)

tensor([[-3.3823e-01, -1.7512e+01,  4.0542e-19,  7.6920e-18],
        [ 5.0572e-02, -3.3823e-01, -7.0284e-20,  1.3939e-19],
        [-1.4957e-19, -8.2174e-18, -2.9764e-01, -6.0422e+00],
        [ 6.4815e-20, -4.3138e-19,  1.5084e-01, -2.9764e-01]], dtype=torch.float64)

tensor([[-3.3823e-01, -1.7512e+01,  4.0542e-19,  7.6920e-18],
        [ 5.0572e-02, -3.3823e-01, -7.0284e-20,  1.3939e-19],
        [-1.4957e-19, -8.2174e-18, -2.9764e-01, -6.0422e+00],
        [ 6.4815e-20, -4.3138e-19,  1.5084e-01, -2.9764e-01]], dtype=torch.float64)

tensor([-1.8863e-15, -1.8863e-15, -2.1015e-16, -2.1015e-16], dtype=torch.float64)

[7]:
# Locate period five fixed points

# Set fixed point period

power = 5

# Set tolerance epsilon

epsilon = 1.0E-9

# Set random initial points

qx = 1.0*torch.rand(256, dtype=torch.float64) - 0.50
px = 0.1*torch.rand(256, dtype=torch.float64) - 0.05
qy = torch.zeros_like(qx)
py = torch.zeros_like(px)

points = torch.stack([qx, px, qy, py]).T

# Perform root search iterations for each initial point

ring, *_ = group(RING, 'FODO_A', 'FODO_D', ('ms', ['Sextupole'], None, None), ('dp', None, None, None), root=True, alignment=False)
points = torch.func.vmap(lambda guess: orbit(ring, guess, ms, dp, limit=128, power=power, epsilon=None))(points)

# Iterate

for _ in range(128):
    locals = torch.vmap(lambda state: ring(state, ms, dp))(points)

# Remove solutions with large norms

points = points[locals.norm(1, dim=-1) < 0.5]

# Remove unconverged solutions

mask = []
for point in points:
    local = point.clone()
    for _ in range(power):
        local = ring(local, ms, dp)
    mask.append((local - point).norm() < epsilon)

points = points[mask]

# Clean points (remove nans, duplicates, points from the same chain)

points = clean_point(power, ring, points, ms, dp, epsilon=epsilon)

# Generate fixed point chains

chains = torch.func.vmap(lambda point: chain_point(power, ring, point, ms, dp))(points)

# Classify fixed point chains (elliptic vs hyperbolic)
# Generate initials for hyperbolic fixed points using corresponding eigenvectors

kinds = []
for chain in chains:
    point, *_ = chain
    values, vectors = torch.linalg.eig(matrix(power, ring, point, ms, dp))
    kind = values.log().real.prod() < epsilon
    kinds.append(bool(kind))
    if not kind:
        lines = [point + vector*torch.linspace(-epsilon, +epsilon, 128, dtype=torch.float64).reshape(-1, 1) for vector in vectors.real.T]
        lines = torch.stack(lines).reshape(-1, 4)

# Remove vertical plane in chains

qx, px, *_ = chains.swapaxes(0, -1)
chains = torch.stack([qx, px]).swapaxes(0, -1)

# Iterate lines and remove vertical plane

manifold = []
for _ in range(64):
    manifold.append(lines)
    lines = torch.func.vmap(lambda point: ring(point, ms, dp))(lines)
manifold = torch.stack(manifold)

# Remove vertical plane in lines (including nonlinear leaking)

qx, px, qy, py = manifold.swapaxes(0, -1)
qx = qx[qy.abs() + py.abs() < epsilon]
px = px[qy.abs() + py.abs() < epsilon]
manifold = torch.stack([qx, px])

# Plot

plt.figure(figsize=(6, 6))
qx, px, *_ = torch.stack(table).swapaxes(0, -1)
plt.scatter(qx.cpu().numpy(), px.cpu().numpy(), s=1, color='black')
qx, px = manifold
plt.scatter(qx.flatten().cpu().numpy(), px.flatten().cpu().numpy(), s=1, color='grey', alpha=0.5)
for chain, kind in zip(chains, kinds):
    plt.scatter(*chain.T, color = {True:'blue', False:'red'}[kind], marker='o')
plt.xlim(-1.0, 0.5)
plt.ylim(-0.075, 0.075)
plt.tight_layout()
plt.show()
../_images/examples_model_388_0.png

Example-28: Orbit (effect of transverse shift)

[1]:
# In this example effect of transverse shifts on closed orbit is illustrated
# Quadrupoles and sextupoles are shifted in transverse planes
# Given shifts variations, effect on orbit at a singe location is computed with MC sampling
[2]:
# Import

import torch
torch.set_printoptions(linewidth=128)

import matplotlib
from matplotlib import pyplot as plt
matplotlib.rcParams['text.usetex'] = True

from twiss import twiss

from ndmap.signature import chop
from ndmap.evaluate import evaluate
from ndmap.pfp import parametric_fixed_point

from ndmap.pfp import clean_point
from ndmap.pfp import chain_point
from ndmap.pfp import matrix

from model.library.drift import Drift
from model.library.quadrupole import Quadrupole
from model.library.sextupole import Sextupole
from model.library.dipole import Dipole
from model.library.line import Line

from model.command.wrapper import wrapper
from model.command.wrapper import group

from model.command.orbit import orbit
[3]:
# Define simple FODO based lattice using nested lines

DR = Drift('DR', 0.25)
BM = Dipole('BM', 3.50, torch.pi/4.0)

QF_A = Quadrupole('QF_A', 0.5, +0.20)
QD_A = Quadrupole('QD_A', 0.5, -0.19)
QF_B = Quadrupole('QF_B', 0.5, +0.20)
QD_B = Quadrupole('QD_B', 0.5, -0.19)
QF_C = Quadrupole('QF_C', 0.5, +0.20)
QD_C = Quadrupole('QD_C', 0.5, -0.19)
QF_D = Quadrupole('QF_D', 0.5, +0.20)
QD_D = Quadrupole('QD_D', 0.5, -0.19)

SF_A = Sextupole('SF_A', 0.25, +0.75)
SD_A = Sextupole('SD_A', 0.25, -1.25)
SF_B = Sextupole('SF_B', 0.25, +0.75)
SD_B = Sextupole('SD_B', 0.25, -1.25)
SF_C = Sextupole('SF_C', 0.25, +0.75)
SD_C = Sextupole('SD_C', 0.25, -1.25)
SF_D = Sextupole('SF_D', 0.25, +0.75)
SD_D = Sextupole('SD_D', 0.25, -1.25)

FODO_A = Line('FODO_A', [QF_A, DR, SF_A, DR, BM, DR, SD_A, DR, QD_A, QD_A, DR, SD_A, DR, BM, DR, SF_A, DR, QF_A], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_B = Line('FODO_B', [QF_B, DR, SF_B, DR, BM, DR, SD_B, DR, QD_B, QD_B, DR, SD_B, DR, BM, DR, SF_B, DR, QF_B], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_C = Line('FODO_C', [QF_C, DR, SF_C, DR, BM, DR, SD_C, DR, QD_C, QD_C, DR, SD_C, DR, BM, DR, SF_C, DR, QF_C], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
FODO_D = Line('FODO_D', [QF_D, DR, SF_D, DR, BM, DR, SD_D, DR, QD_D, QD_D, DR, SD_D, DR, BM, DR, SF_D, DR, QF_D], propagate=True, dp=0.0, exact=False, output=False, matrix=False)

RING = Line('RING', [FODO_A, FODO_B, FODO_C, FODO_D], propagate=True, dp=0.0, exact=False, output=False, matrix=False)
[4]:
# Set transverse error shifts for quadrupoles and sextupoles

ring, table, _ = group(RING,
                       'FODO_A',
                       'FODO_D',
                       ('dx', ['Quadrupole'], None, None),
                       ('dy', ['Quadrupole'], None, None),
                       ('dx', ['Sextupole'], None, None),
                       ('dy', ['Sextupole'], None, None),
                       alignment=True,
                       root=True)

dx_quad = 100.0E-6*torch.randn(8, dtype=torch.float64)
dy_quad = 100.0E-6*torch.randn(8, dtype=torch.float64)

dx_sext = 200.0E-6*torch.randn(8, dtype=torch.float64)
dy_sext = 200.0E-6*torch.randn(8, dtype=torch.float64)

# Test

state = torch.tensor(4*[0.0], dtype=torch.float64)
print(ring(state, dx_quad, dy_quad, dx_sext, dy_sext))
tensor([2.1768e-04, 1.3470e-05, 3.1293e-04, 3.8868e-05], dtype=torch.float64)
[5]:
# Compute and test closed orbit

guess = torch.tensor(4*[0.0], dtype=torch.float64)

fp = orbit(ring, guess,  dx_quad, dy_quad, dx_sext, dy_sext, limit=8, epsilon=1.0E-12)

print(fp)
print(ring(fp, dx_quad, dy_quad, dx_sext, dy_sext))
tensor([2.0328e-05, 1.0872e-05, 6.5715e-05, 3.7633e-05], dtype=torch.float64)
tensor([2.0328e-05, 1.0872e-05, 6.5715e-05, 3.7633e-05], dtype=torch.float64)
[6]:
# Deviation data generation from wrapper

_ = wrapper(RING, *table, verbose=True, alignment=True)
_, data = _(fp, dx_quad, dy_quad, dx_sext, dy_sext)

print(fp)
print(RING(fp, data=data, alignment=True))
tensor([2.0328e-05, 1.0872e-05, 6.5715e-05, 3.7633e-05], dtype=torch.float64)
tensor([2.0328e-05, 1.0872e-05, 6.5715e-05, 3.7633e-05], dtype=torch.float64)
[7]:
# Orbit sensitivity with MC

def fn(dx_quad, dy_quad, dx_sext, dy_sext):
    guess = torch.tensor(4*[0.0], dtype=torch.float64)
    return orbit(ring, guess, dx_quad, dy_quad, dx_sext, dy_sext, limit=64, epsilon=None)

dx_quad = 100.0E-6*torch.randn((2**12, 8), dtype=torch.float64)
dy_quad = 100.0E-6*torch.randn((2**12, 8), dtype=torch.float64)

dx_sext = 200.0E-6*torch.randn((2**12, 8), dtype=torch.float64)
dy_sext = 200.0E-6*torch.randn((2**12, 8), dtype=torch.float64)

cqx, cpx, cqy, cpy = torch.vmap(fn)(dx_quad, dy_quad, dx_sext, dy_sext).T

fig, (ax, ay) = plt.subplots(1, 2, figsize=(12, 5))
ax.hist(cqx.cpu().numpy(), bins=100, range=(-1.0E-3, +1.0E-3), color='blue', alpha=0.7)
ay.hist(cqy.cpu().numpy(), bins=100, range=(-1.0E-3, +1.0E-3), color='blue', alpha=0.7)
plt.tight_layout()
plt.show()
../_images/examples_model_396_0.png
[ ]: